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):