From 7f41613a7e7bc43cb963566d3987d78c172a5110 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 29 Aug 2026 04:38:17 -0700 Subject: [PATCH 1/3] Save the etrecord the lowered paths already generate `--generate_etrecord` produced no file and no warning on the Core ML and XNNPACK llama paths. The flag was not being ignored, which is the part worth knowing: both helpers set `generate_etrecord` on the builder, `to_edge_transform_and_lower` builds the record from it, and the record travels all the way to the ExecuTorch program. Nothing ever saved it. So the cost was already being paid, including a deepcopy of the edge program, and the artifact was dropped at the end. One helper now writes it, from the same `export_program` the combined path uses and to the same `etrecord.bin`, so all three paths leave the same artifact. A missing record is the normal case and stays silent. Also removed the TODO asking for exactly this. One difference worth stating: the record from these paths is about twice the size of the combined path's, because `to_edge_transform_and_lower` also records the aten exported program, which the combined path does not. Measured on the default llama config, 1.65 GB against 826 MB for an 826 MB `.pte`. That is content, not waste, but it is a size a user will notice. Test plan: Two tests driving the real XNNPACK helper, so they fail on the missing file rather than on a missing symbol: base FAILED, AssertionError: Lists differ: [] != ['etrecord.bin'] head 2 passed Also ran the full llama export on the default config with only this file swapped: base ['m.pte'] head ['etrecord.bin', 'm.pte'] and confirmed the written record loads: `parse_etrecord` returns an ETRecord with its edge dialect program set. --- examples/models/llama/export_llama_lib.py | 27 ++++++- .../llama/tests/test_export_llama_lib.py | 74 +++++++++++++++++++ 2 files changed, 98 insertions(+), 3 deletions(-) diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index c22bdac4ed2..e6afc442816 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 @@ -1209,9 +1208,29 @@ def _to_edge_and_lower_llama_xnnpack( print_delegation_info(builder.edge_manager.exported_program().graph_module) # Add gen_tag_fn to tag non-delegated weights as well. - return builder.to_executorch( + builder = builder.to_executorch( passes=additional_passes, external_constants_tag=gen_tag_fn ) + _save_etrecord_if_generated(builder) + return builder + + +def _save_etrecord_if_generated(builder) -> None: + """Write the etrecord the lowering attached, if there is one. + + `to_edge_transform_and_lower` builds the record when asked and carries it through to the + ExecuTorch program, but nothing saves it, so a caller passing --generate_etrecord got no file + and no warning. The combined lowering path calls `generate_etrecord` itself and writes the same + filename, so both paths now leave the same artifact. + """ + try: + etrecord = builder.export_program.get_etrecord() + except RuntimeError: + # Not generated, which is the normal case. + return + + etrecord.save("etrecord.bin") + logging.info("Generated etrecord.bin") def _to_edge_and_lower_llama_openvino( @@ -1392,7 +1411,9 @@ def _to_edge_and_lower_llama_coreml( if verbose: print_delegation_info(builder.edge_manager.exported_program().graph_module) - return builder.to_executorch(passes=additional_passes) + builder = builder.to_executorch(passes=additional_passes) + _save_etrecord_if_generated(builder) + return builder def _to_edge_and_lower_llama( # noqa: C901 diff --git a/examples/models/llama/tests/test_export_llama_lib.py b/examples/models/llama/tests/test_export_llama_lib.py index 40ba701f84e..4c5633a95ae 100644 --- a/examples/models/llama/tests/test_export_llama_lib.py +++ b/examples/models/llama/tests/test_export_llama_lib.py @@ -5,6 +5,8 @@ # 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 executorch.devtools.backend_debug import get_delegation_info @@ -25,6 +27,7 @@ from executorch.examples.models.llama.export_llama_lib import ( _export_llama, + _to_edge_and_lower_llama_xnnpack, build_args_parser, get_quantizer_and_quant_params, ) @@ -40,7 +43,78 @@ ] +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. + """ + import torch + from executorch.extension.llm.export.builder import LLMEdgeManager + + 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 test_xnnpack_lowering_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 the helpers using that call never saved it, so `--generate_etrecord` produced + no file and no warning while still paying for the record. + + Driving the real helper rather than the save step, so this fails on the missing file rather + than on a missing symbol. + """ + builder = _tiny_llm_builder() + + with tempfile.TemporaryDirectory() as directory: + previous = os.getcwd() + os.chdir(directory) + try: + _to_edge_and_lower_llama_xnnpack( + builder, + modelname="tiny", + additional_passes=[], + pt2e_quant_params=None, + quantizers=[], + quant_dtype=None, + generate_etrecord=True, + ) + self.assertEqual(os.listdir("."), ["etrecord.bin"]) + finally: + os.chdir(previous) + + def test_xnnpack_lowering_writes_nothing_when_not_asked(self): + """The common case must stay silent rather than raise or leave a stray file.""" + builder = _tiny_llm_builder() + + with tempfile.TemporaryDirectory() as directory: + previous = os.getcwd() + os.chdir(directory) + try: + _to_edge_and_lower_llama_xnnpack( + builder, + modelname="tiny", + additional_passes=[], + pt2e_quant_params=None, + quantizers=[], + quant_dtype=None, + generate_etrecord=False, + ) + self.assertEqual(os.listdir("."), []) + finally: + os.chdir(previous) + def test_has_expected_ops_and_op_counts(self): """ Checks the presence of unwanted expensive ops. From d361bd71bbfba12344e2ab997acb595d53af4b19 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 30 Aug 2026 10:08:56 -0700 Subject: [PATCH 2/3] Write the record after the model, and never let it lose the export Two problems with the first version, both found in review. The record was written inside the lowering, before the model was saved, and the write was not guarded. So a failed record write took the whole export with it. Measured with an unwritable target: the export raised and left no .pte, where the same run on the previous revision produced one. A record about twice the size of the model makes a full disk the likely trigger, and losing a model to a debug flag is the wrong trade. The record is now written once, after `save_to_pte`, and a write error is logged and swallowed. That also removes the duplicate call from the two lowering helpers. The multimethod path dropped the record too. It builds one exactly as the other paths did and saved only the .pte, so it gets the same call. Test plan: Three tests driving the real export with a stubbed builder: saves the record base ['tiny.pte'], head ['etrecord.bin', 'tiny.pte'] writes none when unasked ['tiny.pte'] keeps the model when the record cannot be written .pte present, warning logged The first fails on the previous revision, so it pins the fix rather than the helper. Also corrected the docstring, which described the state before the change in the present tense and claimed both paths leave the same artifact. They do not: this one carries the aten exported program as well, so it is roughly twice the size. --- examples/models/llama/export_llama_lib.py | 24 ++++--- .../llama/tests/test_export_llama_lib.py | 69 ++++++++++--------- 2 files changed, 51 insertions(+), 42 deletions(-) diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index e6afc442816..5f737af1268 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -1208,20 +1208,17 @@ def _to_edge_and_lower_llama_xnnpack( print_delegation_info(builder.edge_manager.exported_program().graph_module) # Add gen_tag_fn to tag non-delegated weights as well. - builder = builder.to_executorch( + return builder.to_executorch( passes=additional_passes, external_constants_tag=gen_tag_fn ) - _save_etrecord_if_generated(builder) - return builder def _save_etrecord_if_generated(builder) -> None: """Write the etrecord the lowering attached, if there is one. - `to_edge_transform_and_lower` builds the record when asked and carries it through to the - ExecuTorch program, but nothing saves it, so a caller passing --generate_etrecord got no file - and no warning. The combined lowering path calls `generate_etrecord` itself and writes the same - filename, so both paths now leave the same artifact. + 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() @@ -1229,7 +1226,12 @@ def _save_etrecord_if_generated(builder) -> None: # Not generated, which is the normal case. return - etrecord.save("etrecord.bin") + try: + etrecord.save("etrecord.bin") + except OSError as error: + logging.warning("Could not write etrecord.bin: %s", error) + return + logging.info("Generated etrecord.bin") @@ -1411,9 +1413,7 @@ def _to_edge_and_lower_llama_coreml( if verbose: print_delegation_info(builder.edge_manager.exported_program().graph_module) - builder = builder.to_executorch(passes=additional_passes) - _save_etrecord_if_generated(builder) - return builder + return builder.to_executorch(passes=additional_passes) def _to_edge_and_lower_llama( # noqa: C901 @@ -1714,6 +1714,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 @@ -1896,6 +1897,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 4c5633a95ae..89f33640f24 100644 --- a/examples/models/llama/tests/test_export_llama_lib.py +++ b/examples/models/llama/tests/test_export_llama_lib.py @@ -8,6 +8,7 @@ import os import tempfile import unittest +from unittest.mock import patch from executorch.devtools.backend_debug import get_delegation_info @@ -27,7 +28,6 @@ from executorch.examples.models.llama.export_llama_lib import ( _export_llama, - _to_edge_and_lower_llama_xnnpack, build_args_parser, get_quantizer_and_quant_params, ) @@ -65,53 +65,60 @@ def forward(self, tokens): class ExportLlamaLibTest(unittest.TestCase): - def test_xnnpack_lowering_saves_the_etrecord_it_generates(self): + def _run_tiny_export(self, generate_etrecord): + llm_config = LlmConfig() + llm_config.backend.xnnpack.enabled = True + llm_config.debug.generate_etrecord = generate_etrecord + # The recipe rejects dynamic shapes, and _validate_args runs before the lowering. + llm_config.model.enable_dynamic_shape = False + 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 the helpers using that call never saved it, so `--generate_etrecord` produced - no file and no warning while still paying for the record. - - Driving the real helper rather than the save step, so this fails on the missing file rather - than on a missing symbol. + program, but nothing saved it, so the flag produced no file and no warning while still + paying for the record. """ - builder = _tiny_llm_builder() - with tempfile.TemporaryDirectory() as directory: previous = os.getcwd() os.chdir(directory) try: - _to_edge_and_lower_llama_xnnpack( - builder, - modelname="tiny", - additional_passes=[], - pt2e_quant_params=None, - quantizers=[], - quant_dtype=None, - generate_etrecord=True, - ) - self.assertEqual(os.listdir("."), ["etrecord.bin"]) + self._run_tiny_export(generate_etrecord=True) + self.assertEqual(sorted(os.listdir(".")), ["etrecord.bin", "tiny.pte"]) finally: os.chdir(previous) - def test_xnnpack_lowering_writes_nothing_when_not_asked(self): + def test_export_writes_no_etrecord_when_not_asked(self): """The common case must stay silent rather than raise or leave a stray file.""" - builder = _tiny_llm_builder() + 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 and its failure is logged, so a full disk or an + unwritable directory loses the record and not the .pte. + """ with tempfile.TemporaryDirectory() as directory: previous = os.getcwd() os.chdir(directory) try: - _to_edge_and_lower_llama_xnnpack( - builder, - modelname="tiny", - additional_passes=[], - pt2e_quant_params=None, - quantizers=[], - quant_dtype=None, - generate_etrecord=False, - ) - self.assertEqual(os.listdir("."), []) + # A directory of that name makes the write fail without touching permissions. + os.mkdir("etrecord.bin") + self._run_tiny_export(generate_etrecord=True) + self.assertIn("tiny.pte", os.listdir(".")) finally: os.chdir(previous) From 6e6b9dbc30878ee87cf495db53bf9722f1fe8887 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 30 Aug 2026 20:58:15 -0700 Subject: [PATCH 3/3] Keep the promise on every path, and make the tests prove the artifact Review found the guard did not hold where it was written down. The docstring said the save is never allowed to raise, but only OSError was caught. Injecting a RuntimeError from the save showed it escaping the export after the .pte was already on disk, which is the exact outcome the change exists to prevent. Widened to Exception, which is what save_pte_program on the line above already does. The combined lowering was the third path and still wrote its record before the model, unguarded, so a failed write there lost the export. That write is older than this change, but it breaks the same rule, so it now warns and continues too. Three test corrections: The positive test only checked the file name, so a zero byte record passed. It now loads the record back and asserts the edge program, which fails when the save is replaced by an empty file. The failure test asserted only that the .pte survived, which is also true on the previous revision where nothing wrote a record. It now asserts the warning it is named for. Its docstring claimed an unwritable directory loses the record and not the .pte. Measured: in a read only directory both are lost and the export still reports success, because the model save swallows its own error. Dropped that half. Also removed a dynamic shape line and its comment from the test setup. Validation rejects dynamic shapes only for Core ML and QNN, and these tests enable XNNPACK, so nothing read the flag. Hoisted torch and the builder import to module scope to match the sibling test files. --- examples/models/llama/export_llama_lib.py | 18 +++++++++------- .../llama/tests/test_export_llama_lib.py | 21 ++++++++++++------- 2 files changed, 25 insertions(+), 14 deletions(-) diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index 5f737af1268..98363769e08 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -1228,7 +1228,7 @@ def _save_etrecord_if_generated(builder) -> None: try: etrecord.save("etrecord.bin") - except OSError as error: + except Exception as error: logging.warning("Could not write etrecord.bin: %s", error) return @@ -1567,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: diff --git a/examples/models/llama/tests/test_export_llama_lib.py b/examples/models/llama/tests/test_export_llama_lib.py index 89f33640f24..c7aa4f5b5c0 100644 --- a/examples/models/llama/tests/test_export_llama_lib.py +++ b/examples/models/llama/tests/test_export_llama_lib.py @@ -10,7 +10,9 @@ 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 ( @@ -31,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, @@ -48,8 +51,6 @@ def _tiny_llm_builder(): Enough to exercise the lowering helpers without a checkpoint or a real llama model. """ - import torch - from executorch.extension.llm.export.builder import LLMEdgeManager class Tiny(torch.nn.Module): def forward(self, tokens): @@ -69,8 +70,6 @@ def _run_tiny_export(self, generate_etrecord): llm_config = LlmConfig() llm_config.backend.xnnpack.enabled = True llm_config.debug.generate_etrecord = generate_etrecord - # The recipe rejects dynamic shapes, and _validate_args runs before the lowering. - llm_config.model.enable_dynamic_shape = False llm_config.export.output_name = "tiny.pte" with patch( "executorch.examples.models.llama.export_llama_lib._prepare_for_llama_export", @@ -91,6 +90,10 @@ def test_export_saves_the_etrecord_it_generates(self): 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) @@ -108,8 +111,8 @@ def test_export_writes_no_etrecord_when_not_asked(self): 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 and its failure is logged, so a full disk or an - unwritable directory loses the record and not the .pte. + 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() @@ -117,8 +120,12 @@ def test_export_keeps_the_model_when_the_etrecord_cannot_be_written(self): try: # A directory of that name makes the write fail without touching permissions. os.mkdir("etrecord.bin") - self._run_tiny_export(generate_etrecord=True) + 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)