OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

[serialization] Add pte file to archive (#162520)

Browse Source 
Summary:
Add _package_executorch_files to archive apis. Allow us to package a PTE file into the archive.

I don't think there's a use-case to have more than one PTE file at the moment, but left it as `EXECUTORCH_FILES` just in case.

Test Plan:
Tested in D81992612

Rollback Plan:

Differential Revision: D81977483

Pull Request resolved: https://github.com/pytorch/pytorch/pull/162520
Approved by: https://github.com/angelayi
This commit is contained in:
Lucy Qiu 2025-09-11 07:59:08 +00:00 committed by PyTorch MergeBot
parent 6b9b7ce6fe
commit 22df9332da
4 changed files with 20 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
torch/_C/_export/pt2_archive_constants.pyi
View File

@ -18,6 +18,7 @@ TENSOR_CONSTANT_FILENAME_PREFIX: str = ...
CUSTOM_OBJ_FILENAME_PREFIX: str = ...
SAMPLE_INPUTS_DIR: str = ...
SAMPLE_INPUTS_FILENAME_FORMAT: str = ...
EXECUTORCH_DIR: str = ...
EXTRA_DIR: str = ...
MODULE_INFO_PATH: str = ...
XL_MODEL_WEIGHTS_DIR: str = ...

2
torch/csrc/export/pt2_archive_constants.h
View File

@ -47,6 +47,8 @@ namespace torch::_export::archive_spec {
DO(SAMPLE_INPUTS_DIR, "data/sample_inputs/") \
DO(SAMPLE_INPUTS_FILENAME_FORMAT, \
"data/sample_inputs/{}.pt") /* {model_name} */ \
/* ExecuTorch artifacts, including PTE files */ \
DO(EXECUTORCH_DIR, "data/executorch/") \
/* extra folder */ \
DO(EXTRA_DIR, "extra/") \
DO(MODULE_INFO_PATH, "extra/module_info.json") \

16
torch/export/pt2_archive/_package.py
View File

@ -43,6 +43,7 @@ from torch.export.pt2_archive.constants import (
CONSTANTS_CONFIG_FILENAME_FORMAT,
CONSTANTS_DIR,
CUSTOM_OBJ_FILENAME_PREFIX,
EXECUTORCH_DIR,
EXTRA_DIR,
MODELS_DIR,
MODELS_FILENAME_FORMAT,
@ -529,6 +530,16 @@ def _package_extra_files(
archive_writer.write_string(f"{EXTRA_DIR}{extra_file_name}", content)
def _package_executorch_files(
archive_writer: PT2ArchiveWriter, executorch_files: Optional[dict[str, bytes]]
) -> None:
if executorch_files is None:
return
for file_name, content in executorch_files.items():
archive_writer.write_bytes(f"{EXECUTORCH_DIR}{file_name}", content)
def package_pt2(
f: FileLike,
*,
@ -539,6 +550,7 @@ def package_pt2(
extra_files: Optional[dict[str, Any]] = None,
opset_version: Optional[dict[str, int]] = None,
pickle_protocol: int = DEFAULT_PICKLE_PROTOCOL,
executorch_files: Optional[dict[str, bytes]] = None,
) -> FileLike:
r"""
Saves the artifacts to a PT2Archive format. The artifact can then be loaded
@ -569,6 +581,9 @@ def package_pt2(
pickle_protocol: can be specified to override the default protocol
executorch_files (Optional[dict[str, bytes]]): Optional executorch
artifacts to save.
"""
assert not (
exported_programs is None and aoti_files is None and extra_files is None
@ -602,6 +617,7 @@ def package_pt2(
pickle_protocol=pickle_protocol,
)
_package_extra_files(archive_writer, extra_files)
_package_executorch_files(archive_writer, executorch_files)
if isinstance(f, (io.IOBase, IO)):
f.seek(0)

1
torch/export/pt2_archive/constants.py
View File

@ -13,6 +13,7 @@ CONSTANTS_CONFIG_FILENAME_FORMAT: str = (
pt2_archive_constants.CONSTANTS_CONFIG_FILENAME_FORMAT
)
CUSTOM_OBJ_FILENAME_PREFIX: str = pt2_archive_constants.CUSTOM_OBJ_FILENAME_PREFIX
EXECUTORCH_DIR: str = pt2_archive_constants.EXECUTORCH_DIR
EXTRA_DIR: str = pt2_archive_constants.EXTRA_DIR
MODELS_DIR: str = pt2_archive_constants.MODELS_DIR
MODELS_FILENAME_FORMAT: str = pt2_archive_constants.MODELS_FILENAME_FORMAT