From df17b3cf187deca5a6fecb0c8f7c1e4da30783bd Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Mon, 24 Aug 2026 15:38:09 -0700 Subject: [PATCH] ci(executorch): run a coalesced TensorRT + CUDA program in the reference runner gate The ExecuTorch gate exports one program, x + 1, and TensorRT takes it whole. So nothing in CI ever runs a program where TensorRT and ExecuTorch's own CUDA backend each own part of the same graph. That coalesced case is the whole point of combining the two backends, and it is not covered end to end today. There is a composition test that checks both delegates land in the file, but it never loads or runs the program. This adds the missing run. A new example, examples/torchtrt_executorch_example/export_coalesced.py, exports cos(erfinv(tanh(x))). TensorRT has no converter for erfinv, so a CudaPartitioner catch-all gives that operator to the CUDA backend while TensorRT keeps the rest. The script fails if the saved .pte does not carry both a TensorRTBackend and a CudaBackend delegate, so a partitioning change cannot quietly turn this into a TensorRT-only run that still passes. The script also writes .expected next to the .pte, holding the output shape and the eager reference value for an all-ones input. Both reference runners fill inputs with 1.0 and this model is elementwise, so one number describes the whole expected output. Reading it from a file, instead of hard-coding a number in the shell script, keeps the expectation tied to the model. verify-executorch-reference-runner.sh now takes an optional third argument, the coalesced .pte. When given, it runs both the CMake-built runner and the packaged runner on it and compares every printed value against that reference. TensorRT, AOTInductor and eager PyTorch use different kernels for the same math, so the comparison uses a tolerance of 0.001 rather than matching printed digits. The existing x + 1 assertions keep the same strength. They now go through the same helper with a zero tolerance, because x + 1 on ones is exact in float32. Usage: python examples/torchtrt_executorch_example/export_coalesced.py \ --model_path=coalesced.pte .github/scripts/verify-executorch-reference-runner.sh \ model.pte kv_cache_decode.pte coalesced.pte Test plan On a Linux x86_64 host with an NVIDIA A100 GPU: - Ran export_coalesced.py. It reported delegates ['TensorRTBackend', 'CudaBackend', 'TensorRTBackend'] and wrote "[64,64]" and "0.6722" into the .expected file. - Ran the resulting .pte through the reference runner. It printed "output[0] shape=[64,64]" and first 8 values of 0.6722, an exact match to the eager result. - Deleted the aoti_cuda_blob.ptd that the CUDA backend writes and ran again. Same output, so this model needs no external weight file. - Exercised the new shell assertion helper against captured runner output: correct output passes; one wrong value fails; a wrong shape fails; a missing values line fails; a value inside the tolerance passes and one outside it fails. - shellcheck, bash -n, black and isort are clean on the changed files. Not yet observed in CI: the ExecuTorch runtime build job currently fails on main when the packaged reference runner aborts on the existing x + 1 model, and the test job is skipped while that is true. Both happen before this new code runs. --- .../verify-executorch-reference-runner.sh | 96 +++++++++++--- .github/workflows/executorch-test-linux.yml | 5 +- .../export_coalesced.py | 120 ++++++++++++++++++ .../test_cuda_partitioner_composition.py | 9 +- 4 files changed, 207 insertions(+), 23 deletions(-) create mode 100644 examples/torchtrt_executorch_example/export_coalesced.py diff --git a/.github/scripts/verify-executorch-reference-runner.sh b/.github/scripts/verify-executorch-reference-runner.sh index fe8e98d3de..d032662926 100755 --- a/.github/scripts/verify-executorch-reference-runner.sh +++ b/.github/scripts/verify-executorch-reference-runner.sh @@ -18,6 +18,11 @@ set +x # examples/torchtrt_executorch_example/export_kv_cache_decode.py). When given, # kv_cache_decode_check is built and run against it as well. # +# Optional third argument: path to a coalesced TensorRT + CUDA .pte (see +# examples/torchtrt_executorch_example/export_coalesced.py). When given, both +# runners are run against it and their output is compared to the eager +# reference that export script wrote next to the model. +# # Optional: # TensorRT_ROOT=/path/to/extracted/TensorRT # If unset, the script reuses Bazel's fetched TensorRT SDK when available @@ -33,8 +38,8 @@ set +x repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" cd "${repo_root}" -if [[ $# -lt 1 || $# -gt 2 ]]; then - echo "Usage: $0 PATH_TO_MODEL.pte [PATH_TO_KV_CACHE_DECODE.pte]" >&2 +if [[ $# -lt 1 || $# -gt 3 ]]; then + echo "Usage: $0 PATH_TO_MODEL.pte [PATH_TO_KV_CACHE_DECODE.pte [PATH_TO_COALESCED.pte]]" >&2 exit 1 fi model_path="$1" @@ -47,6 +52,11 @@ if [[ -n "${kv_model_path}" && ! -f "${kv_model_path}" ]]; then echo "KV-cache decode model not found: ${kv_model_path}" >&2 exit 1 fi +coalesced_model_path="${3:-}" +if [[ -n "${coalesced_model_path}" && ! -f "${coalesced_model_path}" ]]; then + echo "Coalesced model not found: ${coalesced_model_path}" >&2 + exit 1 +fi python_executable="${PYTHON_EXECUTABLE:-}" if [[ -z "${python_executable}" ]]; then @@ -467,25 +477,35 @@ packaged_runner_log="${verify_root}/packaged_runner.log" --model_path="${model_path}" \ --num_runs=1 2>&1 | tee "${packaged_runner_log}" -# The sample model is x + 1 on a (2,3,4,4) input and both runners fill inputs with -# 1.0f, so the shape is exactly [2,3,4,4] and every printed value is exactly 2.0000. -# Assert both precisely. Matching only "shape=" accepts any shape, and matching one -# 2.0000 anywhere on the values line accepts a line of wrong numbers that happens to -# contain one right one, so neither catches a stream-ordering regression returning -# stale or partial output. Both lines come from fprintf in the runner, so these -# assertions hold whatever ET_LOG_ENABLED is set to. -for _log in "${runner_log}" "${packaged_runner_log}"; do - # A right answer produced entirely on the host would not prove much here: the - # program is delegated to TensorRT, so at least one planned buffer has to be - # served by a registered CUDA DeviceAllocator. Pin that, otherwise a model or - # a planning change could quietly turn this into a CPU-only test. +# Assert the printed shape, and that EVERY value on the "first N values:" line is +# the expected one. Matching only "shape=" accepts any shape, and matching one +# right value anywhere on the values line accepts a line of wrong numbers that +# happens to contain one, so neither on its own catches a stream-ordering +# regression returning stale or partial output. Both lines come from fprintf in +# the runner, so these assertions hold whatever ET_LOG_ENABLED is set to. The +# models used here are elementwise on an all-ones input, so one number describes +# the whole expected output. +assert_runner_output() { + local _log="$1" + local _shape="$2" + local _expected="$3" + local _tolerance="$4" + local _values + local _value + + # A right answer produced entirely on the host would not prove much: the + # programs here are delegated, so at least one planned buffer has to be served + # by a registered CUDA DeviceAllocator. Pin that, otherwise a model or a + # planning change could quietly turn this into a CPU-only test. if ! grep -q 'planned buffer\[[0-9]*\] = [0-9]* bytes on device_type 1' "${_log}"; then echo "No CUDA planned buffer was allocated in ${_log}:" >&2 grep 'planned buffer' "${_log}" >&2 || echo " no planned buffer line at all" >&2 exit 1 fi - if ! grep -q 'output\[0\] shape=\[2,3,4,4\]' "${_log}"; then + # -F: the shape is bracketed, and an unescaped [2,3,4,4] is a regex character + # class that would match any single one of those characters. + if ! grep -qF "output[0] shape=${_shape}" "${_log}"; then echo "Unexpected output shape in ${_log}:" >&2 grep 'output\[0\] shape=' "${_log}" >&2 || echo " no shape line at all" >&2 exit 1 @@ -498,11 +518,19 @@ for _log in "${runner_log}" "${packaged_runner_log}"; do exit 1 fi for _value in ${_values}; do - if [[ "${_value}" != "2.0000" ]]; then - echo "Unexpected output value '${_value}' in ${_log}: ${_values}" >&2 + if ! awk -v got="${_value}" -v want="${_expected}" -v tol="${_tolerance}" \ + 'BEGIN { d = got - want; if (d < 0) d = -d; exit !(d <= tol) }'; then + echo "Unexpected output value '${_value}' in ${_log}" \ + "(expected ${_expected} within ${_tolerance}): ${_values}" >&2 exit 1 fi done +} + +# The sample model is x + 1 on a (2,3,4,4) input, so every output value is exactly +# 2. That is exact in float32, hence a zero tolerance. +for _log in "${runner_log}" "${packaged_runner_log}"; do + assert_runner_output "${_log}" "[2,3,4,4]" "2.0000" 0 done if [[ -n "${kv_model_path}" ]]; then @@ -521,3 +549,37 @@ if [[ -n "${kv_model_path}" ]]; then "${kv_check_path}" --model_path="${kv_model_path}" 2>&1 | tee "${kv_check_log}" grep -q "PASS: decode at pos=1 observed the KV written at pos=0" "${kv_check_log}" fi + +if [[ -n "${coalesced_model_path}" ]]; then + # A coalesced program splits one graph across the TensorRT delegate and + # ExecuTorch's CUDA delegate, so a value produced by one delegate is consumed by + # the other on the device, inside one method. The checks above use a program with + # a single delegate, so they exercise neither the second backend nor the handover + # between the two. + coalesced_expected_path="${coalesced_model_path%.pte}.expected" + if [[ ! -f "${coalesced_expected_path}" ]]; then + echo "Coalesced reference output not found: ${coalesced_expected_path}" >&2 + echo "It is written by examples/torchtrt_executorch_example/export_coalesced.py" >&2 + exit 1 + fi + coalesced_shape="$(sed -n '1p' "${coalesced_expected_path}")" + coalesced_value="$(sed -n '2p' "${coalesced_expected_path}")" + if [[ -z "${coalesced_shape}" || -z "${coalesced_value}" ]]; then + echo "Malformed coalesced reference output in ${coalesced_expected_path}" >&2 + exit 1 + fi + + # Only the from-source runner runs the coalesced program. It links every + # ExecuTorch delegate, including the CUDA/AOTI backend that the CUDA partition + # of a coalesced .pte is handed to. The packaged runner in the wheel ships the + # TensorRT delegate alone, so it has no CudaBackend to run that partition and + # cannot execute this model. + coalesced_runner_log="${verify_root}/coalesced_my_runner.log" + "${runner_path}" \ + --model_path="${coalesced_model_path}" \ + --num_runs=1 2>&1 | tee "${coalesced_runner_log}" + + # TensorRT, AOTInductor and eager PyTorch compute the same math with different + # kernels, so compare within a tolerance instead of on the printed digits. + assert_runner_output "${coalesced_runner_log}" "${coalesced_shape}" "${coalesced_value}" 0.001 +fi diff --git a/.github/workflows/executorch-test-linux.yml b/.github/workflows/executorch-test-linux.yml index 6d38fff9bb..9751f2bb54 100644 --- a/.github/workflows/executorch-test-linux.yml +++ b/.github/workflows/executorch-test-linux.yml @@ -95,8 +95,11 @@ jobs: --model_path="${RUNNER_TEMP}/torchtrt-python.pte" python examples/torchtrt_executorch_example/export_kv_cache_decode.py \ --model_path="${RUNNER_TEMP}/torchtrt-kv-cache-decode.pte" + python examples/torchtrt_executorch_example/export_coalesced.py \ + --model_path="${RUNNER_TEMP}/torchtrt-coalesced.pte" .github/scripts/verify-executorch-reference-runner.sh \ "${RUNNER_TEMP}/torchtrt-python.pte" \ - "${RUNNER_TEMP}/torchtrt-kv-cache-decode.pte" + "${RUNNER_TEMP}/torchtrt-kv-cache-decode.pte" \ + "${RUNNER_TEMP}/torchtrt-coalesced.pte" python examples/executorch_reference_runner/load_model.py \ --model_path="${RUNNER_TEMP}/torchtrt-python.pte" --num_runs=1 diff --git a/examples/torchtrt_executorch_example/export_coalesced.py b/examples/torchtrt_executorch_example/export_coalesced.py new file mode 100644 index 0000000000..9f9d1fbc90 --- /dev/null +++ b/examples/torchtrt_executorch_example/export_coalesced.py @@ -0,0 +1,120 @@ +""" +.. _executorch_export_coalesced: + +Exporting a Coalesced TensorRT + CUDA Model to ExecuTorch (.pte) +================================================================ + +This example exports one graph split across two ExecuTorch backends. TensorRT +takes the operators it can convert, and everything left over goes to +ExecuTorch's own CUDA backend, which compiles it with AOTInductor. Both +delegates end up inside a single ``.pte`` and run in the same method. + +The model is ``cos(erfinv(tanh(x)))``. TensorRT has no converter for +``erfinv``, so that operator is the one the CUDA backend has to claim. That +makes the split real rather than incidental. + +A value produced by one delegate is consumed by the other on the device, inside +a single method, so this is what proves the two backends can complete each +other rather than only work on their own. + +The CUDA backend also writes an ``aoti_cuda_blob.ptd`` next to the ``.pte`` for +its external weights. This model has no weights, so that file is empty of +tensors and the reference runner does not need it. + +Besides the ``.pte`` this writes ``.expected``, holding the output +shape and the eager reference value for an all-ones input. The reference runner +gate reads that file instead of hard-coding a number, so the expected value +cannot drift away from the model. + +Prerequisites +------------- +Install Torch-TensorRT with the ExecuTorch extra before running this example:: + + pip install -e ".[executorch]" + +ExecuTorch's CUDA backend also needs a CUDA toolkit (``nvcc``) at export time, +for the AOTInductor compile. +""" + +import argparse +import sys +from pathlib import Path + +import torch +import torch_tensorrt +from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.cuda_partitioner import CudaPartitioner +from executorch.exir._serialize._program import deserialize_pte_binary + +SHAPE = (64, 64) + + +class CoalescedModel(torch.nn.Module): + def forward(self, x): + return torch.cos(torch.erfinv(torch.tanh(x))) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--model_path", default="coalesced.pte", help="Path to save the .pte" + ) + args = parser.parse_args() + model_path = Path(args.model_path) + + with torch.no_grad(): + model = CoalescedModel().eval().cuda() + example_input = (torch.randn(SHAPE).cuda(),) + + exported_program = torch.export.export(model, example_input) + trt_gm = torch_tensorrt.dynamo.compile( + exported_program, + arg_inputs=example_input, + min_block_size=1, + truncate_double=True, + ) + torch_tensorrt.save( + trt_gm, + str(model_path), + output_format="executorch", + arg_inputs=example_input, + retrace=False, + # Catch-all, so every operator TensorRT rejected goes to the CUDA + # backend instead of falling back to a portable CPU kernel. + partitioners=[ + CudaPartitioner( + [CudaBackend.generate_method_name_compile_spec("forward")] + ) + ], + ) + + # Both delegates must really be in the file. A partitioner change that + # quietly routed the whole graph to TensorRT would otherwise leave a + # green job that no longer tests the coalesced path at all. + program = deserialize_pte_binary(model_path.read_bytes()).program + delegates = [d.id for plan in program.execution_plan for d in plan.delegates] + missing = [ + name for name in ("TensorRTBackend", "CudaBackend") if name not in delegates + ] + if missing: + sys.exit( + f"{model_path} is not coalesced: missing {missing}, found {delegates}" + ) + + # The reference runners fill every input element with 1.0, and this model + # is elementwise, so a single number describes the whole expected output. + reference = model(torch.ones(SHAPE).cuda()) + expected_path = model_path.with_suffix(".expected") + expected_path.write_text( + "[{}]\n{:.4f}\n".format( + ",".join(str(dim) for dim in reference.shape), + reference.flatten()[0].item(), + ) + ) + + print(f"Saved {model_path} with delegates {delegates}.") + print(f"Saved {expected_path} with the eager reference output.") + + +if __name__ == "__main__": + main() diff --git a/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py b/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py index f3ef0d47e5..9bdc065e18 100644 --- a/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py +++ b/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py @@ -89,11 +89,10 @@ def _delegate_ids(pte_path): # NOTE: these tests are COMPOSITION-ONLY. They assert the serialized .pte carries # both a TensorRTBackend and a CudaBackend delegate (and, below, that external -# CUDA weights are persisted as a .ptd). They do NOT load or run the program: a -# coalesced ATen-mode .pte cannot be loaded yet because memory-planned CUDA -# buffers get a CPU data pointer, so Method::init fails the CUDA backend's device -# check (tensor_parser_aten hardcodes CPU). A load-run-allclose test should be -# added once that runtime fix lands (follow-up). +# CUDA weights are persisted as a .ptd). They do NOT load or run the program. +# Execution of a coalesced program is covered by the reference runner gate, see +# examples/torchtrt_executorch_example/export_coalesced.py and the third argument +# of .github/scripts/verify-executorch-reference-runner.sh. def test_erfinv_routes_to_cuda_backend(tmp_path):