[HiroForYou]

Summary

_get_extended_header() in exir/_serialize/_program.py and extension/flat_tensor/serialize/serialize.py, along with FlatTensorSerializer.deserialize(), read the extended header using from_bytes(data[8:]).

Slicing a bytes object copies the slice, so data[8:] copies almost the entire file just to read a header that is at most 40 bytes.

This change limits the slice to the header bytes (8 : 8 + EXPECTED_LENGTH), which is all from_bytes() reads. Parsing results are unchanged.

For .pte files, the copy happens in _get_extended_header(), which deserialize_pte_binary() calls on the entire file. For .ptd files, it happened in FlatTensorSerializer.deserialize(), which had its own header read.

FlatTensorSerializer.deserialize() and parse_executorch_program() in backends/mlx/pte_inspector.py now use _get_extended_header() instead of reading the header separately. This keeps the header parsing in one place and avoids duplicating the logic. One behavior change: deserialize() now raises RuntimeError ("Flat tensor header is invalid") instead of ValueError when the input is too short. I couldn't find any caller that relies on the ValueError.

I also removed extension/flat_tensor/test/test_serialize.py from pytest.ini's --ignore list so the .ptd regression test runs in the open-source CI job. The file shares its basename with devtools/bundled_program/serialize/test/test_serialize.py. Since neither directory had an __init__.py, a repo-wide pytest run failed during collection with import file mismatch. Adding an empty extension/flat_tensor/test/__init__.py lets pytest import the test as flat_tensor.test.test_serialize, avoiding the clash.

Test plan

Added regression tests that fail on main and pass with this change. Both wrap from_bytes() and check that it receives exactly EXPECTED_LENGTH bytes:

• exir/_serialize/test/test_program.py::test_get_extended_header_reads_only_header_bytes covers the .pte helper, which parse_executorch_program() now uses too.
• extension/flat_tensor/test/test_serialize.py::test_reads_only_header_bytes covers the .ptd helper through FlatTensorSerializer.deserialize().

Reverting either slice to data[8:], data[8:len(data)], or bytes(data)[8:] makes one of these tests fail.

python -m pytest exir/_serialize/test/ extension/flat_tensor/test/  # 105 passed on this branch, 94 on main

lintrunner exir/_serialize/_program.py extension/flat_tensor/serialize/serialize.py \
  exir/_serialize/test/test_program.py extension/flat_tensor/test/test_serialize.py \
  backends/mlx/pte_inspector.py

Tested with a 384 MiB .pte (6 x Linear(4096, 4096), XNNPACK) and a 384 MiB .ptd (same model exported with external_constants=True). Each operation was run 7 times in one process. The reported time is the median of the last 6 calls, then the median across 3 alternating before/after rounds. Memory is the increase in peak RSS of the Python process.

Operation	Before	After
_get_extended_header() on the .pte	307 ms, +384 MiB	<1 ms, +0 MiB
deserialize_pte_binary()	625 ms	291 ms
FlatTensorSerializer().deserialize()	614 ms	291 ms
MLX parse_executorch_program()	1392 ms, +384 MiB	1094 ms, +0 MiB

This avoids a 384 MiB copy in each path, saving around 300 ms in these measurements. Peak memory during full deserialization is unchanged (+384 MiB), since restoring the segments still copies the weights.

For the MLX inspector, this removes a full copy in the Python process. After reading the header, parse_executorch_program() passes the entire file to flatc in a child process, which loads it again and peaks near 770 MiB. As a result, the overall peak memory for this path remains largely unchanged.

In a fresh process, the first call is slower: the old slice also faults in 384 MiB of new pages, which the segment copies then reuse. So on a first call the header looks more expensive on its own than the saving on the full path, which is why the table uses the median of the repeated calls.

The parsed headers, PTEFile, DataPayload, and MLX inspector output produce identical hashes before and after the change.

Scripts to reproduce the measurements

# make_files.py: writes model.pte and _default_external_constant.ptd

import torch

from torch.export import export

from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
from executorch.exir import to_edge_transform_and_lower
from executorch.exir.capture._config import ExecutorchBackendConfig

m = torch.nn.Sequential(
    *[torch.nn.Linear(4096, 4096) for _ in range(6)]
).eval()

x = (torch.randn(1, 4096),)

pte = to_edge_transform_and_lower(
    export(m, x), partitioner=[XnnpackPartitioner()]
)
open("model.pte", "wb").write(pte.to_executorch().buffer)

ext = to_edge_transform_and_lower(export(m, x)).to_executorch(
    ExecutorchBackendConfig(external_constants=True)
)
ext.write_tensor_data_to_file(".")  # writes _default_external_constant.ptd

# bench.py OP FILE
#   python bench.py pte_header model.pte
#   python bench.py pte_deserialize model.pte
#   python bench.py ptd_deserialize _default_external_constant.ptd
#   python bench.py mlx_inspect model.pte  # needs the generated MLX files

import resource
import statistics
import sys
import time

from executorch.exir._serialize import _program
from executorch.exir._serialize._cord import Cord
from executorch.extension.flat_tensor.serialize.serialize import FlatTensorSerializer

op_name, path = sys.argv[1:]
data = open(path, "rb").read()

if op_name == "mlx_inspect":
    from executorch.backends.mlx.pte_inspector import parse_executorch_program

ops = {
    "pte_header": lambda: _program._get_extended_header(data),
    "pte_deserialize": lambda: _program.deserialize_pte_binary(data),
    "ptd_deserialize": lambda: FlatTensorSerializer().deserialize(Cord(data)),
    "mlx_inspect": lambda: parse_executorch_program(data),
}

base = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
times = []

for _ in range(7):
    start = time.perf_counter()
    ops[op_name]()
    times.append((time.perf_counter() - start) * 1000)

mib = (resource.getrusage(resource.RUSAGE_SELF).ru_maxrss - base) / 1024

print(
    f"first call {times[0]:.0f} ms, then median {statistics.median(times[1:]):.0f} ms,"
    f" peak +{mib:.0f} MiB"
)

Authored with assistance from Claude Code.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23561

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 5 New Failures, 11 Unrelated Failures

As of commit 789d217 with merge base 0295293 ():

NEW FAILURES - The following jobs have failed:

• Test QNN Windows MSVC build / build-qnn-windows-arm64 (gh)
• trunk / test-models-macos-coreml (dl3) / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• trunk / test-models-macos-coreml (ic4) / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• trunk / test-models-macos-coreml (mobilebert) / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• trunk / test-models-macos-coreml (resnet50) / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• trunk / test-models-macos-coreml (mv2) / macos-job (gh) (similar failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• trunk / test-models-macos-coreml (vit) / macos-job (gh) (similar failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• trunk / test-models-macos-coreml (w2l) / macos-job (gh) (similar failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Apple / build-demo-ios / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
• trunk / test-coreml-delegate / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
• trunk / test-models-macos-coreml (edsr) / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• trunk / test-models-macos-coreml (efficient_sam) / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• trunk / test-models-macos-coreml (emformer_join) / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• trunk / test-models-macos-coreml (emformer_transcribe) / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• trunk / test-models-macos-coreml (ic3) / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• trunk / test-models-macos-coreml (mv3) / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[HiroForYou]

@pytorchbot label "release notes: none"

[HiroForYou]

@shoumikhin since you worked on this function in #23340, could you take a look when you get a chance?

[shoumikhin]

Looks good, thanks. backends/mlx/pte_inspector.py makes the same whole-file copy when it reads the header, with pte_data[extended_header_offset:]. You could fix that here too, if you like.

[HiroForYou]

@shoumikhin I applied the same fix to backends/mlx/pte_inspector.py as you suggested. On the 384 MiB .pte, parse_executorch_program() peak memory drops from +384 MiB to +0 MiB, with identical output. I didn't add a test because the MLX tests require generated MLX files.

[shoumikhin]

Four notes, all small. The slicing change itself is right, and the .ptd test reads well.

Two of them are about the numbers in the description rather than the code.

The MLX row reads as a peak memory win, and for the Python process it is. But right after the header, parse_executorch_program() hands the whole file to _program_flatbuffer_to_json(), which writes it to a temp file and runs flatc, and flatc loads the file again in a child process. On a 384 MiB input that child peaks near 770 MiB either way, so the peak for the run as a whole does not really move. "One less full copy in Python" might describe it better.

The other is the two .pte rows. deserialize_pte_binary() calls _get_extended_header() on the whole file, so the saving on the outer call should be at least the saving on the inner one. The table has 455 ms for the header on its own and 198 ms for the full path. Allocator state in a fresh process could account for some of that, since the isolated benchmark has to fault in 384 MiB of fresh pages. Could you say which, or re-measure? bench.py covers the header row and the .ptd row but not the deserialize_pte_binary() row or the MLX one, so it is hard to check from the outside.

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[HiroForYou]

@shoumikhin, thanks for the second pass :). You were right about both numbers, and I updated the description:

• MLX: It now says the change removes one full copy in the Python process. I measured flatc at about 770 MiB with the change, so the overall peak memory for the run doesn't really change.
• The two .pte rows: The difference came from allocator state. In a fresh process, the old header copy also faults in 384 MiB of new pages, which the segment copies can then reuse. That made the header look more expensive on its own than the actual savings in the full path. bench.py now covers all four rows and reports the first call separately from the median of repeated calls. With repeated calls, each path saves about one 384 MiB copy (~300 ms), including the header-only call.

I also merged main into the branch. The Apple failure was caused by a 404 for the cached torch wheel, which #23543 fixed on main.

[nil-is-all]

@HiroForYou could you update the branch? Seems to be out-of-date.

[HiroForYou]

@HiroForYou could you update the branch? Seems to be out-of-date.

@nil-is-all done, the branch is up to date with main

[shoumikhin]

Looks good to me. The header read is bounded in one place now, deserialize() goes through the same helper instead of repeating it, and the empty __init__.py is a smaller fix than the rename. Thanks for measuring the import path question rather than taking my word for it.

Two things before this goes in.

The workflows on this push are still waiting to be approved and run, so the jobs that went red last time have not run again yet. Please get those green, since that was the whole reason for the previous round.

A few lines in the description no longer match the code:

• The part about dropping the ignore does not mention the __init__.py that makes dropping it safe. On its own the removal is what broke the jobs, so the two belong together.
• It still says that reverting any of the three slices makes a test fail. There are two slice sites now, which matches the six cases you reported in the thread rather than nine.
• The test plan labels 105 passed as the number on main. That command gives 94 on main, and 105 is this branch.

[HiroForYou]

@shoumikhin thanks for the review and all the suggestions :)

I fixed the test count in the description, it's 105 on this branch and 94 on main, as you pointed out. The other two points should already be fixed.

I also merged the latest main into the branch. The workflows are waiting for approval.

[HiroForYou]

@nil-is-all @shoumikhin I updated the branch with the latest changes from main since the previous CI runs had expired and the PR was blocked. The new workflows are waiting for approval. Once they pass, it's ready to merge. Thanks!