From 36054d172266d35c71780b4412e02f129c9abc62 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 11 Oct 2026 16:47:08 -0700 Subject: [PATCH] Save a lowered program whose named data is file-backed A backend can hand its named data to ExecuTorch as FileBackedData instead of bytes, so a large blob waits on disk rather than as a second copy in memory. The CUDA backend stores its weights this way. NamedDataStore, to_executorch() and write_tensor_data_to_file() all handle it. exir.save() of a program that contains such a lowered module fails, though: TypeError: Object of type FileBackedData is not JSON serializable The serializer turns the lowered module's named data into JSON with an encoder that knows only bytes. Loading already reads every buffer back as bytes. The encoder for named data now reads a file-backed buffer into bytes when it reaches it, and then encodes it like any other buffer. The lowered module keeps its file-backed buffers, and the saved file is the same as if the backend had given bytes. Like any buffer in this JSON format, the data is held in memory while it is encoded, a few times over. The named data of a lowered module is still written once for each call to it, as it already was for bytes. Tested with a new test in exir/tests/test_serde.py. It lowers a small graph with the named-data test backend, which gives it two buffers, and makes one of them a FileBackedData that owns its file, as the CUDA backend's are. It checks that the saved program is the same as when both buffers are bytes, that the module and its file still hold the data after saving, and that the program loads back with both buffers in order. It fails without this change with the error above. All 16 tests in that file pass with it. --- exir/serde/serialize.py | 11 +++++++++- exir/tests/targets.bzl | 1 + exir/tests/test_serde.py | 44 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 55 insertions(+), 1 deletion(-) diff --git a/exir/serde/serialize.py b/exir/serde/serialize.py index 2df4ba79876..d9cd9a700a7 100644 --- a/exir/serde/serialize.py +++ b/exir/serde/serialize.py @@ -22,6 +22,7 @@ import torch import torch.export.exported_program as ep from executorch.exir import delegate +from executorch.exir._serialize._cord import FileBackedData from executorch.exir._serialize._named_data_store import NamedDataStoreOutput from executorch.exir.backend.backend_details import DelegateScratchSpec from executorch.exir.backend.compile_spec_schema import ( @@ -48,6 +49,14 @@ log: logging.Logger = logging.getLogger(__name__) +class _NamedDataEncoder(export_serialize.EnumEncoder): + # Read each file only when it is encoded, so its raw bytes are freed right after. + def default(self, obj): + if isinstance(obj, FileBackedData): + obj = obj.to_bytes() + return super().default(obj) + + class GraphModuleSerializer(export_serialize.GraphModuleSerializer): def __init__( self, @@ -291,7 +300,7 @@ def serialize_bytes(b: bytes) -> str: assert isinstance(serialized_artifact.exported_program, schema.ExportedProgram) serialized_processed_bytes = serialize_bytes(lowered_module.processed_bytes) - named_data_store = json.dumps(export_serialize._dataclass_to_dict(lowered_module.named_data_store_output),cls=export_serialize.EnumEncoder) if lowered_module.named_data_store_output else None + named_data_store = json.dumps(export_serialize._dataclass_to_dict(lowered_module.named_data_store_output),cls=_NamedDataEncoder) if lowered_module.named_data_store_output else None serialized_lowered_module = SerdeLoweredBackendModule( original_module=serialized_artifact.exported_program, diff --git a/exir/tests/targets.bzl b/exir/tests/targets.bzl index d8f21020d2a..ff161d755e6 100644 --- a/exir/tests/targets.bzl +++ b/exir/tests/targets.bzl @@ -80,6 +80,7 @@ def define_common_targets(is_fbcode = False): "//executorch/exir:lib", "//executorch/exir/backend:backend_api", "//executorch/exir/backend/test:backend_with_compiler_demo", + "//executorch/exir/backend/test:backend_with_named_data_map", "//executorch/exir/backend/test:op_partitioner_demo", "//executorch/exir/serde:serialize", "//executorch/backends/xnnpack/partition:xnnpack_partitioner", diff --git a/exir/tests/test_serde.py b/exir/tests/test_serde.py index 6696b95fd31..2a159c704f8 100644 --- a/exir/tests/test_serde.py +++ b/exir/tests/test_serde.py @@ -7,6 +7,7 @@ # pyre-strict import io +import os import tempfile import unittest from contextlib import contextmanager @@ -19,12 +20,17 @@ XnnpackFloatingPointPartitioner, ) from executorch.exir import to_edge +from executorch.exir._serialize._cord import FileBackedData from executorch.exir.backend.backend_api import CompileSpec, to_backend from executorch.exir.backend.test.backend_with_compiler_demo import ( BackendWithCompilerDemo, ) +from executorch.exir.backend.test.backend_with_named_data_map import ( + BackendWithNDMPartitioner, +) from executorch.exir.backend.test.op_partitioner_demo import AddMulPartitionerDemo +from executorch.exir.lowered_backend_module import get_lowered_submodules from executorch.exir.program._program import ( EdgeProgramManager, to_edge_transform_and_lower, @@ -246,6 +252,44 @@ def forward(self, a, x, b): edge_new = deserialize(serialize(edge.exported_program())) self.check_ep(edge.exported_program(), edge_new, inputs) + def test_delegate_with_file_backed_named_data(self) -> None: + class Model(nn.Module): + def forward(self, x): + return (x + x) - x + + edge = to_edge(export(Model(), (torch.randn(1, 2),), strict=True)).to_backend( + BackendWithNDMPartitioner() + ) + graph_module = edge.exported_program().graph_module + (_, lowered, _) = get_lowered_submodules(graph_module)[0] + output = lowered.named_data_store_output + assert output is not None + payload = bytes(range(256)) + output.buffers[1] = payload + other = output.buffers[0] + expected = serialize(edge.exported_program()).exported_program + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "data.bin") + with open(path, "wb") as f: + f.write(payload) + data = FileBackedData.move_from(path) + output.buffers[1] = data + + serialized = serialize(edge.exported_program()).exported_program + self.assertEqual(serialized, expected) + buffer = io.BytesIO() + exir.save(edge.exported_program(), buffer) + self.assertIs(lowered.named_data_store_output, output) + self.assertIs(output.buffers[1], data) + self.assertEqual(data.to_bytes(), payload) + + buffer.seek(0) + loaded = exir.load(buffer) + (_, loaded_lowered, _) = get_lowered_submodules(loaded.graph_module)[0] + loaded_output = loaded_lowered.named_data_store_output + assert loaded_output is not None + self.assertEqual(loaded_output.buffers, [other, payload]) + def test_delegate_xnnpack(self) -> None: class SimpleConv1DModel(nn.Module): def __init__(self):