diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 29cba8b5ada..2e00a64bc24 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -1234,7 +1234,17 @@ class ET_EXPERIMENTAL CudaBackend final const NamedDataMap* named_data_map, const std::string& weights_blob_key) const { auto buffer_res = named_data_map->get_data(weights_blob_key.c_str()); - if (buffer_res.ok() && handle->update_constants_from_blob != nullptr) { + // A library built before external weights exports no bind function, so + // there is nothing to do with a blob whether or not one was supplied. + if (handle->update_constants_from_blob == nullptr) { + ET_LOG( + Info, + "weights_blob '%s' is not used: this library exports no bind function", + weights_blob_key.c_str()); + return Error::Ok; + } + + if (buffer_res.ok()) { ET_LOG(Info, "Found %s in named data map", weights_blob_key.c_str()); const void* weights_blob = buffer_res->data(); auto update_err = handle->update_constants_from_blob( @@ -1246,10 +1256,39 @@ class ET_EXPERIMENTAL CudaBackend final ET_CUDA_CHECK_OR_RETURN_ERROR(cudaDeviceSynchronize()); buffer_res->Free(); } else { + // The container expects a blob and did not get one, so its constant + // pointers stay null and the failure resurfaces much later as an illegal + // access inside a generated kernel. Report it while the cause is still + // identifiable, as the cached path already does. A model with no + // constants needs nothing bound and stays valid. + // Both symbols are optional, and the count one has been present for + // longer, so a library exporting the update function without it is not + // something any released toolchain produces. Refuse rather than guess, + // since guessing wrong either rejects a valid model or loads one with + // unbound constants. + ET_CHECK_OR_RETURN_ERROR( + handle->get_num_constants != nullptr, + NotSupported, + "weights_blob '%s' is unavailable and this library cannot report its constant count, so whether it needs the blob cannot be established", + weights_blob_key.c_str()); + size_t num_constants = 0; + ET_CHECK_OK_OR_RETURN_ERROR( + handle->get_num_constants(handle->container_handle, &num_constants), + "Failed to enumerate CUDA AOTI constants"); + ET_CHECK_OR_RETURN_ERROR( + num_constants == 0, + NotFound, + "weights_blob '%s' is unavailable (0x%" PRIx32 + "), but the model has %zu constant(s) to bind", + weights_blob_key.c_str(), + static_cast(buffer_res.error()), + num_constants); ET_LOG( Info, - "weights_blob '%s' not found or update fn is null", - weights_blob_key.c_str()); + "weights_blob '%s' is unavailable (0x%" PRIx32 + "), and the model has no constants", + weights_blob_key.c_str(), + static_cast(buffer_res.error())); } return Error::Ok; } diff --git a/backends/cuda/tests/test_missing_weights_blob.py b/backends/cuda/tests/test_missing_weights_blob.py new file mode 100644 index 00000000000..d96d4c70d04 --- /dev/null +++ b/backends/cuda/tests/test_missing_weights_blob.py @@ -0,0 +1,171 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""A model whose weights blob is absent must fail to load, not fail later. + +The CUDA backend hands a model's constants to its generated library through a +sidecar blob. When that blob is missing the container keeps null constant +pointers, and the load used to succeed: the failure surfaced from the first +execute as an illegal memory access inside a generated kernel, with nothing in +the message naming the blob. + +Two payload shapes reach that code. A current export carries per-name weight +metadata and goes through the weight cache, which reports a missing blob itself. +A library built before external weights carries only the two blob keys, newline +separated, and goes through the legacy path. That legacy path is the one the +check was added to, so it is the one this file exercises, by rewriting the +payload of a real export into the older shape in place. +""" + +import os +import tempfile +import unittest + +import torch +from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.cuda_partitioner import CudaPartitioner + +from executorch.backends.cuda.cuda_weight_collector import CUDA_WEIGHT_CACHE_MAGIC +from executorch.exir import to_edge_transform_and_lower +from executorch.exir._serialize._program import deserialize_pte_binary +from torch.export import export + + +class ModuleWithConstants(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("a", torch.randn(4, 64)) + self.register_buffer("b", torch.randn(64)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = x * self.a + y = y + self.b + y = torch.relu(y) + return y * self.a + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") +@unittest.skipIf( + torch.version.hip is not None, + "loads through the Python runtime, which the ROCm build does not include", +) +class TestMissingWeightsBlob(unittest.TestCase): + def _lower(self, outdir: str) -> str: + torch.manual_seed(0) + module = ModuleWithConstants().eval().cuda() + inputs = (torch.randn(4, 64, device="cuda"),) + exported = export(module, inputs) + spec = CudaBackend.generate_method_name_compile_spec("forward") + lowered = to_edge_transform_and_lower( + exported, partitioner=[CudaPartitioner([spec])] + ) + program = lowered.to_executorch() + path = os.path.join(outdir, "model.pte") + with open(path, "wb") as f: + program.write_to_file(f) + program.write_tensor_data_to_file(outdir) + return path + + def _rewrite_payload_as_legacy(self, path: str) -> str: + """Replace the weight metadata payload with the older two-key payload. + + Rewriting in place keeps every offset in the file valid, and modelling the + older shape this way avoids needing an old toolchain to produce one. The + payload is read from the program rather than located by scanning, because a + program with more than one delegate holds more than one payload and a scan + cannot tell where one ends. + """ + with open(path, "rb") as f: + raw = f.read() + program = deserialize_pte_binary(raw).program + + payloads = [] + for plan in program.execution_plan: + for delegate in plan.delegates: + if delegate.id != "CudaBackend": + continue + stored = program.backend_delegate_data[delegate.processed.index].data + payloads.append( + b"".join(stored.chunks()) + if hasattr(stored, "chunks") + else bytes(stored) + ) + self.assertEqual(len(payloads), 1, "expected one CUDA delegate") + payload = payloads[0] + self.assertTrue( + payload.startswith(CUDA_WEIGHT_CACHE_MAGIC), + "expected the weight metadata payload this rewrite consumes", + ) + + # The payload carries a content hash, so it occurs once. + offset = raw.find(payload) + self.assertGreaterEqual(offset, 0, "payload not found in the file") + self.assertEqual( + raw.find(payload, offset + 1), -1, "payload is not unique in the file" + ) + + marker = b"_so_blob" + end = payload.index(marker) + len(marker) + so_key = payload[:end].rsplit(b"\x00", 1)[-1].decode("utf-8") + weights_key = so_key.replace("_so_blob", "_weights_blob") + + # The two keys, then zeros to keep the payload its original length. The + # shared library key ends at the newline and stays exact; the filler lands + # in the blob key, which is still a key the data map does not hold. + legacy = (so_key + "\n" + weights_key).encode("utf-8") + self.assertLessEqual(len(legacy), len(payload), "legacy payload does not fit") + legacy += b"\x00" * (len(payload) - len(legacy)) + + blob = bytearray(raw) + blob[offset : offset + len(payload)] = legacy + with open(path, "wb") as f: + f.write(bytes(blob)) + return weights_key + + def test_load_reports_not_found_when_blob_is_absent(self) -> None: + from executorch.runtime import Runtime + + with tempfile.TemporaryDirectory() as outdir: + path = self._lower(outdir) + + blobs = [f for f in os.listdir(outdir) if f.endswith(".ptd")] + self.assertTrue(blobs, "expected an externalized weights blob") + # A blob holding no constants is a bare header, and a missing blob + # really is harmless then, so the model has to carry real data for + # this to be testing anything. + largest = max(os.path.getsize(os.path.join(outdir, f)) for f in blobs) + self.assertGreater(largest, 256, "expected non-empty constants") + + # A positive control first: the same program loads when its blob is + # supplied, so a later failure is about the missing blob and not about + # the program, the backend registration or this rewrite. + sidecar = os.path.join(outdir, blobs[0]) + Runtime.get().load_program(path, data_path=sidecar).load_method("forward") + + self._rewrite_payload_as_legacy(path) + + # Without this the test would still pass if the rewrite stopped + # working, by exercising the weight cache path instead, which reports + # the same error number for the same program. + with open(path, "rb") as f: + self.assertNotIn( + CUDA_WEIGHT_CACHE_MAGIC, + f.read(), + "the rewrite left the metadata payload in place", + ) + + # The blob is never supplied, so the load must fail. The runtime's + # exception carries only the method name and the error number, so the + # cause is asserted through the rewrite check above rather than here. + program = Runtime.get().load_program(path) + with self.assertRaisesRegex( + RuntimeError, r"Failed to load method forward, error: 0x:?20" + ): + program.load_method("forward") + + +if __name__ == "__main__": + unittest.main()