diff --git a/.ci/scripts/wheel/test_cuda_linux.py b/.ci/scripts/wheel/test_cuda_linux.py index 5a8e20b1665..914fafaa3f3 100644 --- a/.ci/scripts/wheel/test_cuda_linux.py +++ b/.ci/scripts/wheel/test_cuda_linux.py @@ -320,11 +320,15 @@ def test_the_delegate_registers() -> None: # search path is in place before glibc caches it. Kept as source rather than a separate file # because the smoke test is invoked directly and ships as one module. _EXECUTION_CHILD = """ +import os +import tempfile + import torch from executorch.backends.cuda.cuda_partitioner import CudaPartitioner from executorch.exir import to_edge_transform_and_lower from executorch.extension.pybindings.portable_lib import ( + _load_for_executorch, _load_for_executorch_from_buffer, ) @@ -357,6 +361,99 @@ def forward(self, x, y): torch.testing.assert_close(actual.cpu(), eager, rtol=1e-3, atol=1e-3) capability = torch.cuda.get_device_capability(0) print(f"PASS: a CUDA-delegated model ran on sm_{capability[0]}{capability[1]} and matched eager") + + +# The model above has no weights, so it needs nothing from the data file and runs whether or not +# the delegate's external data path works. Anything a user would actually ship has weights, and +# the CUDA delegate does not put them in the .pte: it writes them to a separate aoti_cuda_blob.ptd +# that the caller has to supply at load time. Nothing here loaded one, so a wheel that shipped a +# broken data path passed this check. Measured on sm_80: given the same model without the blob, +# the delegate reports the weights as not found, and the run then dies inside forward with an +# illegal memory access rather than returning wrong numbers. +class Weighted(torch.nn.Module): + def __init__(self): + super().__init__() + self.fc = torch.nn.Linear(8, 4) + + def forward(self, x): + return self.fc(x) + + +weighted = Weighted().eval() +weighted_example = (torch.randn(2, 8),) +with torch.no_grad(): + weighted_eager = weighted(*weighted_example) + +weighted_program = to_edge_transform_and_lower( + torch.export.export(weighted, weighted_example), partitioner=[CudaPartitioner([])] +).to_executorch() + +assert b"CudaBackend" in weighted_program.buffer, ( + "the weighted program carries no CUDA delegate, so it would run on the portable kernels and " + "never reach the external weights path this check exists for" +) + +# Only the weight matrix, not every parameter. Measured on sm_80, the delegate wrote the 128-byte +# weight of this Linear into the data file and kept its 16-byte bias out, so the sum over +# parameters() is not a bound on what the file has to hold. +weight_bytes = weighted.fc.weight.numel() * weighted.fc.weight.element_size() + +with tempfile.TemporaryDirectory() as work_dir: + pte_path = os.path.join(work_dir, "weighted.pte") + with open(pte_path, "wb") as handle: + weighted_program.write_to_file(handle) + weighted_program.write_tensor_data_to_file(work_dir) + + blob_path = os.path.join(work_dir, "aoti_cuda_blob.ptd") + assert os.path.exists(blob_path), ( + "the export wrote no aoti_cuda_blob.ptd, so the weights of this model went nowhere and " + f"every runner script that passes that file by name would fail: {sorted(os.listdir(work_dir))}" + ) + + # A data file is written even for a model with no weights, and it is not small: measured on + # sm_80, the weightless model above produces a 256-byte container around an empty payload. So + # the file's presence says nothing and its raw size says nothing either. What means something + # is the amount by which this file exceeds that empty container, so write one and subtract it + # instead of comparing against a constant, which would stop catching anything the next time + # the container grows. + empty_dir = os.path.join(work_dir, "weightless") + os.mkdir(empty_dir) + lowered.write_tensor_data_to_file(empty_dir) + empty_size = os.path.getsize(os.path.join(empty_dir, "aoti_cuda_blob.ptd")) + + blob_size = os.path.getsize(blob_path) + assert blob_size - empty_size >= weight_bytes, ( + f"aoti_cuda_blob.ptd is {blob_size} bytes against {empty_size} bytes for a model with no " + f"weights, a difference of {blob_size - empty_size}, which does not cover this model's " + f"{weight_bytes}-byte weight matrix, so the weights were not written to it" + ) + + with open(pte_path, "rb") as handle: + pte_bytes = handle.read() + with open(blob_path, "rb") as handle: + blob_bytes = handle.read() + + # The runner scripts pass this file by path, not as bytes, so exercise the + # path-based loader here while the file still exists, or a regression in it + # stays green. + file_module = _load_for_executorch(pte_path, blob_path) + file_actual = file_module.forward(list(weighted_example))[0] + torch.testing.assert_close( + file_actual.cpu(), weighted_eager, rtol=1e-3, atol=1e-3 + ) + +# Loading from bytes as well, because _load_for_executorch_from_buffer is what a +# caller with the data already in memory uses, and it builds the data map by a +# different route. +module = _load_for_executorch_from_buffer(pte_bytes, blob_bytes) +actual = module.forward(list(weighted_example))[0] + +torch.testing.assert_close(actual.cpu(), weighted_eager, rtol=1e-3, atol=1e-3) +print( + f"PASS: a CUDA-delegated model with weights ran on sm_{capability[0]}{capability[1]} from a " + f"{blob_size}-byte aoti_cuda_blob.ptd carrying {blob_size - empty_size} bytes of weights, " + f"loaded both by path and from bytes, and matched eager" +) """ @@ -396,12 +493,17 @@ def _glibcxx_versions(library: Path) -> Set[str]: def test_a_model_runs_through_the_delegate() -> None: - """Export a model to the CUDA delegate and run it, comparing against eager. + """Export two models to the CUDA delegate and run them, comparing against eager. Every other check in this file reads a shipped artifact. This one executes, because a wheel whose libraries are all present and correctly linked can still fail to compute, and nothing above would notice. The x86_64 rows land on a GPU runner, so this is the row where that can be proven. + The second model carries weights. The CUDA delegate keeps those out of the .pte and writes them + to a separate aoti_cuda_blob.ptd, which the caller then has to supply, so a wheel can run the + weightless model and still be unable to run anything real. That path is only covered by the + second model. + Only the aarch64 rows may skip, and they say why, so a green result never stands for work that did not happen. On x86_64 the absent device or the mismatched torch build is itself the failure: that row is the one place execution is proven, and letting it report success for a check it did