From 768f56ee9a0ae9ba954584454e04e78cd2b6868b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 31 Aug 2026 08:04:57 -0700 Subject: [PATCH 1/6] 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. A shared helper now writes it, from the same `export_program` and to the same `etrecord.bin` name the combined path uses. 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 larger than the combined path's, because `to_edge_transform_and_lower` also records the aten exported program, which the combined path does not. Roughly twice the size on the same model. That is content, not waste, but it is a size a user will notice. Test plan: Tests driving the real export rather than the save helper, so they fail on the missing file rather than on a missing symbol: base the record is absent head the record is written beside the model 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 5cdc9a40940..b878f6b5fcd 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 f33f4489c91..651ca1cc693 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 unittest.mock import patch @@ -27,6 +29,7 @@ from executorch.examples.models.llama import export_llama_lib 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, ) @@ -42,6 +45,27 @@ ] +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 _assert_routes_to(self, lowering, coreml=False, vulkan=False, qnn=False): """Assert which lowering an export routes to, without running one.""" @@ -96,6 +120,56 @@ 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 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 06d2475605c08b11823ca53b8039ad680b15f42d Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 31 Aug 2026 08:04:57 -0700 Subject: [PATCH 2/6] 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 the 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 | 68 ++++++++++--------- 2 files changed, 50 insertions(+), 42 deletions(-) diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index b878f6b5fcd..b4671a6eee4 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 @@ -1894,6 +1895,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 651ca1cc693..ed19b805ac4 100644 --- a/examples/models/llama/tests/test_export_llama_lib.py +++ b/examples/models/llama/tests/test_export_llama_lib.py @@ -29,7 +29,6 @@ from executorch.examples.models.llama import export_llama_lib 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, ) @@ -120,53 +119,60 @@ 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 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 2039cff9b567cb60c4e51a3a337497d851192fd1 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 31 Aug 2026 08:04:57 -0700 Subject: [PATCH 3/6] 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 b4671a6eee4..6f8848e6cd5 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 ed19b805ac4..0d91a9b5282 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 ( @@ -32,6 +34,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, @@ -49,8 +52,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): @@ -123,8 +124,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", @@ -145,6 +144,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) @@ -162,8 +165,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() @@ -171,8 +174,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) From 4d8a357401bfa9c2b0dcabca98d2b1bd147d3cac Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 31 Aug 2026 08:32:49 -0700 Subject: [PATCH 4/6] Write the record where the model goes, and stop exporting twice Review found the record ignored the output directory. The model honoured `--output-dir` and the record did not, so with that flag set the larger of the two artifacts landed in the shell's working directory, and a script collecting the requested directory missed it. Both write sites now join the builder's output directory, so the two cannot disagree. I had argued this was better left alone because the combined path already wrote a bare name. That was the wrong call: the builder carries the directory and both sites already had it in hand, so fixing both is smaller than explaining why neither is fixed. Also removed a redundant export from the test helper. `_export_llama` exports the builder it is given, so returning an already exported one traced the model twice per test. --- examples/models/llama/export_llama_lib.py | 15 +++++++++------ .../models/llama/tests/test_export_llama_lib.py | 10 ++++------ 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index 6f8848e6cd5..0f151197d49 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -14,6 +14,7 @@ import json import logging import math +import os import re import shlex from functools import partial @@ -1226,13 +1227,14 @@ def _save_etrecord_if_generated(builder) -> None: # Not generated, which is the normal case. return + path = os.path.join(builder.output_dir, "etrecord.bin") try: - etrecord.save("etrecord.bin") + etrecord.save(path) except Exception as error: - logging.warning("Could not write etrecord.bin: %s", error) + logging.warning("Could not write %s: %s", path, error) return - logging.info("Generated etrecord.bin") + logging.info("Generated %s", path) def _to_edge_and_lower_llama_openvino( @@ -1567,16 +1569,17 @@ def _to_edge_and_lower_llama( # noqa: C901 # Generate ETRecord if edge_manager_copy: + et_record_path = os.path.join(builder.output_dir, "etrecord.bin") try: generate_etrecord_func( - et_record="etrecord.bin", + et_record=et_record_path, edge_dialect_program=edge_manager_copy, executorch_program=builder.export_program, ) except Exception as error: - logging.warning("Could not write etrecord.bin: %s", error) + logging.warning("Could not write %s: %s", et_record_path, error) else: - logging.info("Generated etrecord.bin") + logging.info("Generated %s", et_record_path) 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 0d91a9b5282..a0a86ef837c 100644 --- a/examples/models/llama/tests/test_export_llama_lib.py +++ b/examples/models/llama/tests/test_export_llama_lib.py @@ -48,9 +48,9 @@ def _tiny_llm_builder(): - """An exported LLMEdgeManager small enough to lower in a unit test. + """An LLMEdgeManager small enough to lower in a unit test. - Enough to exercise the lowering helpers without a checkpoint or a real llama model. + `_export_llama` exports it, so this returns it unexported. """ class Tiny(torch.nn.Module): @@ -63,7 +63,7 @@ def forward(self, tokens): max_seq_len=4, use_kv_cache=False, example_inputs=(torch.ones(1, 4, dtype=torch.long),), - ).export() + ) class ExportLlamaLibTest(unittest.TestCase): @@ -177,9 +177,7 @@ def test_export_keeps_the_model_when_the_etrecord_cannot_be_written(self): 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) - ) + self.assertTrue(any("Could not write" in line for line in logs.output)) finally: os.chdir(previous) From 7efeda00672fa670d622435643bb20df6e96691e Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 1 Sep 2026 14:40:31 -0700 Subject: [PATCH 5/6] Write the record beside the model, and never destroy the one already there Three problems with the previous revision, all found in review. The record was written under `export.output_dir`, but an output name ending in `.pte` is used as the path verbatim and so bypasses that directory. With both set, which is the documented form in this repo, the model and its record landed in different directories. Measured on the combined lowering: before this stack both files landed together in the working directory, and after it the model was in the working directory and the record was in the output directory. The record path is now derived from the path the model was actually written to. The record was written in place. The format truncates its target when it opens it and closes in a finally, so a failure part way through left a short file under the real name and destroyed any record already there, reporting only a warning. Worse, the leftover still loaded: the parser accepted a 35479 byte remnant of a 36011 byte record. The record is now staged under a temporary name in the same directory and moved into place, so a failed write leaves the previous record untouched. The combined lowering still wrote its own record before the model, so a full disk could take the model and leave only the record. It now hands the edge program to the same post-model step, which leaves one write site and one rule. Also drop the function-local `import os` that the module-level one added here made dead. Test plan: The three earlier tests pinned none of the three fixes they were named for, because they patched out the only caller of `set_output_dir`, so the builder's output directory was always the default and the join could not be observed. They now point the builder at a temporary directory instead of changing the process working directory, which both makes the directory observable and removes a process-wide `chdir` from a suite that runs under `pytest -n auto`. Five tests now, each checked by reverting the behaviour it is named for: record path back to output_dir -> 1 failure guard narrowed from Exception to OSError -> 1 error staged write replaced by in-place write -> 1 failure record call removed from multimethod -> 1 failure Full file: 16 tests, 10 pass and 6 skip for a backend that is not installed here. Linux x86-64, Python 3.12. --- examples/models/llama/export_llama_lib.py | 74 ++++++---- .../llama/tests/test_export_llama_lib.py | 135 ++++++++++++------ 2 files changed, 145 insertions(+), 64 deletions(-) diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index 0f151197d49..c43619ce60f 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -17,6 +17,7 @@ import os import re import shlex +import tempfile from functools import partial from importlib import resources as _resources from json import JSONDecodeError @@ -763,7 +764,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 @@ -1214,22 +1214,57 @@ 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 attached, if there is one. + """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`. - 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. + 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 + edge_program = getattr(builder, "etrecord_edge_program", None) + if edge_program is None: + try: + etrecord = builder.export_program.get_etrecord() + except RuntimeError: + # Not generated, which is the normal case. + return + else: + etrecord = None - path = os.path.join(builder.output_dir, "etrecord.bin") + path = _etrecord_path_beside_model(builder) try: - etrecord.save(path) + with tempfile.NamedTemporaryFile( + dir=os.path.dirname(path) or ".", suffix=".bin", delete=False + ) as temporary: + staged = temporary.name + try: + if etrecord is not None: + etrecord.save(staged) + else: + generate_etrecord_func( + et_record=staged, + edge_dialect_program=edge_program, + executorch_program=builder.export_program, + ) + os.replace(staged, path) + except BaseException: + os.unlink(staged) + raise except Exception as error: logging.warning("Could not write %s: %s", path, error) return @@ -1567,19 +1602,10 @@ def _to_edge_and_lower_llama( # noqa: C901 passes=additional_passes, ) - # Generate ETRecord + # The record is written after the model, by _save_etrecord_if_generated, so a failed + # record write cannot cost the export and the record lands beside the .pte. if edge_manager_copy: - et_record_path = os.path.join(builder.output_dir, "etrecord.bin") - try: - generate_etrecord_func( - et_record=et_record_path, - edge_dialect_program=edge_manager_copy, - executorch_program=builder.export_program, - ) - except Exception as error: - logging.warning("Could not write %s: %s", et_record_path, error) - else: - logging.info("Generated %s", et_record_path) + builder.etrecord_edge_program = edge_manager_copy 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 a0a86ef837c..8df78e9a5f4 100644 --- a/examples/models/llama/tests/test_export_llama_lib.py +++ b/examples/models/llama/tests/test_export_llama_lib.py @@ -13,6 +13,7 @@ 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 ( @@ -37,6 +38,7 @@ from executorch.extension.llm.export.builder import LLMEdgeManager from executorch.extension.llm.export.config.llm_config import ( LlmConfig, + MethodConfig, Pt2eQuantize, VgfQuantizeScope, ) @@ -120,66 +122,119 @@ 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): + def _run_tiny_export( + self, generate_etrecord, directory, output_name="tiny.pte", output_dir=None + ): + """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. + """ llm_config = LlmConfig() llm_config.backend.xnnpack.enabled = True llm_config.debug.generate_etrecord = generate_etrecord - llm_config.export.output_name = "tiny.pte" + 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=_tiny_llm_builder(), + return_value=builder, ): _export_llama(llm_config) + return sorted(os.listdir(directory)) - def test_export_saves_the_etrecord_it_generates(self): - """A lowering that generates an etrecord must also write it. + def test_export_saves_the_etrecord_beside_the_model(self): + """The record must be written, and land beside the model. - `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. + 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: - 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) + 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: - previous = os.getcwd() - os.chdir(directory) - try: - self._run_tiny_export(generate_etrecord=False) - self.assertEqual(os.listdir("."), ["tiny.pte"]) - finally: - os.chdir(previous) + 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 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. + 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: - 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" in line for line in logs.output)) - finally: - os.chdir(previous) + 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): """ From fba1fb7d48ad610474f4f1e7fb907b584d557f59 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 1 Sep 2026 16:28:55 -0700 Subject: [PATCH 6/6] Take the record from the program on every path The combined lowering was the one path still building the record by hand, with the standalone generator and a deep copy of the whole edge manager. Every other path asks the program for the record the export attached and saves that. The plumbing for the same thing was already there, just not connected: `to_edge` takes the flag, `to_backend` and `to_executorch` both carry the record forward, and `get_etrecord` returns it at the end. The two shared export helpers between the builder and `to_edge` did not forward the flag, so this adds it to both, defaulted off, which leaves their other callers unchanged. That makes the save helper one branch instead of two and removes the deep copy along with the attribute the lowering used to hand the record over. One thing worth knowing about the resulting file. Asking the program for the record gets the aten exported program as well, which the hand-built one did not carry, so the record from this path is now about twice the size for the same model. That is more to debug with, and it matches what the other paths already produced. Test plan: A new test drives the combined lowering, which nothing covered before, and checks the record lands beside the model and loads back. It is not decoration: the first version of this change set the flag at the lowering rather than before the edge export, and that path silently produced no record at all. Removing that one line now fails the test. Seventeen tests pass. Measured on a stack of eight linear layers, peak Python memory is the same either way, 30.8 MiB, so the deep copy was not costing what its comment implied and this is about having one mechanism rather than about memory. The record from that model goes from 8.5 MB to 17 MB, which is the aten program being included. --- examples/models/llama/export_llama_lib.py | 36 ++++++------------- .../llama/tests/test_export_llama_lib.py | 32 +++++++++++++++-- extension/export_util/utils.py | 9 ++++- extension/llm/export/builder.py | 1 + 4 files changed, 49 insertions(+), 29 deletions(-) diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index c43619ce60f..ffe9a89a570 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -26,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, ) @@ -1236,15 +1235,11 @@ def _save_etrecord_if_generated(builder) -> None: 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. """ - edge_program = getattr(builder, "etrecord_edge_program", None) - if edge_program is None: - try: - etrecord = builder.export_program.get_etrecord() - except RuntimeError: - # Not generated, which is the normal case. - return - else: - etrecord = None + try: + etrecord = builder.export_program.get_etrecord() + except RuntimeError: + # Not generated, which is the normal case. + return path = _etrecord_path_beside_model(builder) try: @@ -1253,14 +1248,7 @@ def _save_etrecord_if_generated(builder) -> None: ) as temporary: staged = temporary.name try: - if etrecord is not None: - etrecord.save(staged) - else: - generate_etrecord_func( - et_record=staged, - edge_dialect_program=edge_program, - executorch_program=builder.export_program, - ) + etrecord.save(staged) os.replace(staged, path) except BaseException: os.unlink(staged) @@ -1478,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() @@ -1586,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) @@ -1601,11 +1588,8 @@ def _to_edge_and_lower_llama( # noqa: C901 builder = builder.to_executorch( passes=additional_passes, ) - - # The record is written after the model, by _save_etrecord_if_generated, so a failed - # record write cannot cost the export and the record lands beside the .pte. - if edge_manager_copy: - builder.etrecord_edge_program = edge_manager_copy + # 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: diff --git a/examples/models/llama/tests/test_export_llama_lib.py b/examples/models/llama/tests/test_export_llama_lib.py index 8df78e9a5f4..96640045e33 100644 --- a/examples/models/llama/tests/test_export_llama_lib.py +++ b/examples/models/llama/tests/test_export_llama_lib.py @@ -123,7 +123,12 @@ def test_core_ml_with_vulkan_routes_to_the_combined_lowering(self): self.assertTrue(target.call_args.kwargs["vulkan"]) def _run_tiny_export( - self, generate_etrecord, directory, output_name="tiny.pte", output_dir=None + self, + generate_etrecord, + directory, + output_name="tiny.pte", + output_dir=None, + xnnpack=True, ): """Export the tiny model, writing the .pte into `directory`. @@ -132,10 +137,14 @@ def _run_tiny_export( `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 = True + 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) @@ -146,6 +155,25 @@ def _run_tiny_export( _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. 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