Skip to content

feat(executorch): forward ExecuTorch lowering kwargs through save(output_format="executorch") #4433

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Conarnar wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(executorch): forward ExecuTorch lowering kwargs through save(output_format="executorch") #4433

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 47 additions & 0 deletions docsrc/user_guide/runtime_performance/saving_models.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,53 @@ must be pointed at those data files to load them.
``.pte`` into the same directory overwrites the blob and the first ``.pte``
will fail to load. Save each coalesced model into its own directory.

**ExecuTorch lowering options**

When ``output_format="executorch"``, ``torch_tensorrt.save`` forwards the following
keyword arguments to ExecuTorch's ``to_edge_transform_and_lower(...)``. They are
only consulted for the ``executorch`` format; passing them with any other
``output_format`` logs a warning and is otherwise ignored.

* ``constant_methods`` — a ``dict`` of extra constant methods to embed in the
``.pte`` (e.g. ``{"get_max_seq_len": 2048}`` for an LLM runner).
* ``transform_passes`` — additional edge-dialect transform passes to run before
lowering.
* ``compile_config`` — an ``EdgeCompileConfig``. When omitted, Torch-TensorRT
supplies a default with ``_check_ir_validity=False`` (the TensorRT
``execute_engine`` placeholder graph does not pass edge-IR validation). A
caller-supplied config is forwarded **verbatim**, so if you pass your own and
your graph carries TensorRT engines, set ``_check_ir_validity=False`` explicitly.
* ``backend_config`` — an ``ExecutorchBackendConfig`` forwarded to
``to_executorch(...)``.
* ``generate_etrecord`` — a ``bool`` (default ``False``). When ``True``, an
`ETRecord <https://pytorch.org/executorch/stable/etrecord.html>`_ is written
next to the ``.pte`` as ``<base>_etrecord.bin`` (e.g. ``trt.pte`` →
``trt_etrecord.bin``) for use with the ExecuTorch Developer Tools ``Inspector``.

.. code-block:: python

from executorch.exir import EdgeCompileConfig

torch_tensorrt.save(
trt_gm, "trt.pte", output_format="executorch",
retrace=False, arg_inputs=inputs,
constant_methods={"get_max_seq_len": 2048},
compile_config=EdgeCompileConfig(_check_ir_validity=False),
generate_etrecord=True,
)

The ETRecord sidecar can be parsed back and paired with a runtime ETDump in the
Developer Tools ``Inspector``:

.. code-block:: python

from executorch.devtools import Inspector
from executorch.devtools.etrecord import parse_etrecord

etrecord = parse_etrecord("trt_etrecord.bin")
inspector = Inspector(etdump_path="etdump.etdp", etrecord=etrecord)
inspector.print_data_tabular()


Saving torch.compile models
-----------------------------
Expand Down
71 changes: 69 additions & 2 deletions py/torch_tensorrt/_compile.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -743,7 +743,9 @@ def save(
- If both dynamic_shapes and Input objects are provided, the explicit dynamic_shapes
parameter takes precedence.
kwargs: Additional format-specific kwargs. ``partitioners=``,
``compile_specs=``, and ``backend_config=`` are only used with
``compile_specs=``, ``backend_config=``, ``constant_methods=``,
``transform_passes=``, ``compile_config=`` and
``generate_etrecord=`` are only used with
``output_format="executorch"``; otherwise they are ignored with a
warning. Pass ``compile_specs=[CompileSpec("target_device",
b"cuda:<i>")]`` to override the default target device (``cuda:0``).
Expand All @@ -757,6 +759,17 @@ def save(
<executorch_save>` for the ``CudaPartitioner`` recipe, its
export-time requirements (CUDA backend + nvcc), and the external
``.ptd`` weight caveats.
``constant_methods=``, ``transform_passes=`` and
``compile_config=`` are forwarded to
``to_edge_transform_and_lower(...)``. When ``compile_config`` is
omitted it defaults to ``_check_ir_validity=False`` (the TRT engine
graph fails edge IR validation); a caller-supplied
``compile_config`` is forwarded verbatim, so set
``_check_ir_validity=False`` on it explicitly when the graph carries
TRT engines.
``generate_etrecord=True`` writes an ETRecord to
``<model>_etrecord.bin`` next to the ``.pte`` for use with the
ExecuTorch Inspector.
"""
if isinstance(module, CudaGraphsTorchTensorRTModule):
module = module.compiled_module
Expand Down Expand Up @@ -784,6 +797,10 @@ def save(
executorch_partitioners = kwargs.pop("partitioners", None)
executorch_compile_specs = kwargs.pop("compile_specs", None)
executorch_backend_config = kwargs.pop("backend_config", None)
executorch_constant_methods = kwargs.pop("constant_methods", None)
executorch_transform_passes = kwargs.pop("transform_passes", None)
executorch_compile_config = kwargs.pop("compile_config", None)
executorch_generate_etrecord = kwargs.pop("generate_etrecord", False)

if output_format not in accepted_formats:
raise ValueError(
Expand Down Expand Up @@ -915,6 +932,26 @@ def _extract_tensor(obj: Any) -> Any:
"backend_config= is only used with output_format='executorch' and will "
f"be ignored for output_format='{output_format}'."
)
if executorch_constant_methods and output_format != "executorch":
logger.warning(
"constant_methods= is only used with output_format='executorch' and will "
f"be ignored for output_format='{output_format}'."
)
if executorch_transform_passes and output_format != "executorch":
logger.warning(
"transform_passes= is only used with output_format='executorch' and will "
f"be ignored for output_format='{output_format}'."
)
if executorch_compile_config and output_format != "executorch":
logger.warning(
"compile_config= is only used with output_format='executorch' and will "
f"be ignored for output_format='{output_format}'."
)
if executorch_generate_etrecord and output_format != "executorch":
logger.warning(
"generate_etrecord= is only used with output_format='executorch' and will "
f"be ignored for output_format='{output_format}'."
)
if output_format == "aot_inductor" and platform.system() != "Linux":
raise ValueError(
f"The AOT Inductor format is only supported on Linux, {platform.system()} is not a supported platform for this format"
Expand Down Expand Up @@ -984,6 +1021,10 @@ def _extract_tensor(obj: Any) -> Any:
partitioners=executorch_partitioners,
compile_specs=executorch_compile_specs,
backend_config=executorch_backend_config,
constant_methods=executorch_constant_methods,
transform_passes=executorch_transform_passes,
compile_config=executorch_compile_config,
generate_etrecord=executorch_generate_etrecord,
)
else:
raise RuntimeError(
Expand Down Expand Up @@ -1050,6 +1091,10 @@ def _extract_tensor(obj: Any) -> Any:
partitioners=executorch_partitioners,
compile_specs=executorch_compile_specs,
backend_config=executorch_backend_config,
constant_methods=executorch_constant_methods,
transform_passes=executorch_transform_passes,
compile_config=executorch_compile_config,
generate_etrecord=executorch_generate_etrecord,
)
else:
raise RuntimeError(
Expand Down Expand Up @@ -1137,6 +1182,10 @@ def _extract_tensor(obj: Any) -> Any:
partitioners=executorch_partitioners,
compile_specs=executorch_compile_specs,
backend_config=executorch_backend_config,
constant_methods=executorch_constant_methods,
transform_passes=executorch_transform_passes,
compile_config=executorch_compile_config,
generate_etrecord=executorch_generate_etrecord,
)
else:
raise RuntimeError(
Expand Down Expand Up @@ -1418,15 +1467,33 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None
# would fail trying to cast the CustomObjArgument placeholder to a real C++ Engine object.
exp_program = _replace_execute_engine_for_executorch(exp_program)

# Default to _check_ir_validity=False (get_edge_compile_config): the TRT
# execute_engine placeholder graph fails edge IR validation. A caller who passes
# their own compile_config owns it -- it is forwarded verbatim and we do NOT
# override _check_ir_validity (EdgeCompileConfig defaults it to True, so a caller
# whose graph carries TRT engines should set _check_ir_validity=False explicitly).
compile_config = kwargs.get("compile_config")
if compile_config is None:
compile_config = get_edge_compile_config()

generate_etrecord = kwargs.get("generate_etrecord", False)
edge_program = to_edge_transform_and_lower(
exp_program,
transform_passes=kwargs.get("transform_passes"),
partitioner=partitioners,
compile_config=get_edge_compile_config(),
compile_config=compile_config,
constant_methods=kwargs.get("constant_methods"),
generate_etrecord=generate_etrecord,
)
executorch_program = edge_program.to_executorch(config=kwargs.get("backend_config"))
with open(file_path, "wb") as f:
executorch_program.write_to_file(f)
_write_external_tensor_data(executorch_program, file_path)
if generate_etrecord:
# Follows ExecuTorch's example convention (e.g. examples/cuda/scripts/export.py):
# persist the ETRecord as "<pte_base>_etrecord.bin" next to the .pte.
etrecord_path = os.path.splitext(file_path)[0] + "_etrecord.bin"
executorch_program.get_etrecord().save(etrecord_path)


def _normalize_engine_constants_to_python(exp_program: "ExportedProgram") -> None:
Expand Down
Loading
Loading