Skip to content

Commit 4a61ef4

Browse files
shoumikhinr
andauthored
Cover a CUDA model with weights in the wheel smoke test (pytorch#22089)
### Summary The CUDA rows of the wheel smoke test export one model through the CUDA delegate and run it. That model is `x + y`, so it has no weights. A model with no weights needs nothing from the delegate's external data file, so it runs whether or not that path works. The CUDA delegate does not put weights in the `.pte`. It writes them to a separate file named `aoti_cuda_blob.ptd`, and the caller has to hand that file back at load time. Every runner script in this repository passes it by name. Nothing in the smoke test loaded one, so a wheel that could not produce or consume that file would pass the check and ship. ### What this changes A second model, a small `nn.Linear`, in the same child process: - export it through `CudaPartitioner` and write both the `.pte` and the data file, - assert the program actually carries a CUDA delegate, so a partitioner that silently declined the model cannot pass this check on the portable kernels, - assert the data file is present under the name the runners expect, - assert the data file is larger than the one the weightless model above produces, by at least the size of this model's weight matrix, - load the program with the data file and compare the result against eager. Only the x86_64 rows execute this, the same as the existing check. The aarch64 rows have no GPU on their validation runner and keep skipping with a printed reason. ### Why the size check is written as a difference A data file is written even when the model has no weights, and it is not small. Measured on sm_80, the weightless model above produces a 256-byte file: a container around an empty payload. So any fixed lower bound below 256 is satisfied by a file with nothing in it, and a check like `size >= weight_bytes` for a small model is true no matter what the delegate did. Instead the check writes the weightless program's data file in the same run, which costs nothing because that program is already compiled, and subtracts. That keeps working if the container format changes size later, which a hardcoded 256 would not. Comparing against the sum over `parameters()` would not work either. Measured on sm_80 with this model, the delegate wrote the 128-byte weight into the data file and kept the 16-byte bias out, so that sum is not a bound on what the file has to hold. The check uses the weight matrix only. ### Test plan Ran the new child code on a Linux x86_64 machine with an sm_80 GPU, against a published nightly CUDA wheel. Both models pass: ``` PASS: a CUDA-delegated model ran on sm_80 and matched eager PASS: a CUDA-delegated model with weights ran on sm_80 from a 384-byte aoti_cuda_blob.ptd carrying 128 bytes of weights, and matched eager ``` The delegate log on that run names a real content hash for the weights blob, not the hash of an empty one. Also confirmed the new check is meaningful in both directions. Loading the same weighted program without the data file gives: ``` [cuda_backend.cpp:1242] weights_blob '...' not found or update fn is null [cuda_allocator.cpp:87] cudaMemcpy D2H failed: an illegal memory access was encountered RuntimeError: Failed to execute method forward, error: 0x1 ``` And if the delegate wrote an empty data file for the weighted model, the measured numbers give 256 minus 256, a difference of 0, against a 128-byte weight matrix, so the assertion fires. The previous version of this check did not: it compared the raw file size against the parameter total, and 256 is already larger than 144, so it would have passed on an empty file. --------- Co-authored-by: r <r@e>
1 parent bd1e461 commit 4a61ef4

1 file changed

Lines changed: 103 additions & 1 deletion

File tree

‎.ci/scripts/wheel/test_cuda_linux.py‎

Lines changed: 103 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -320,11 +320,15 @@ def test_the_delegate_registers() -> None:
320320
# search path is in place before glibc caches it. Kept as source rather than a separate file
321321
# because the smoke test is invoked directly and ships as one module.
322322
_EXECUTION_CHILD = """
323+
import os
324+
import tempfile
325+
323326
import torch
324327
325328
from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
326329
from executorch.exir import to_edge_transform_and_lower
327330
from executorch.extension.pybindings.portable_lib import (
331+
_load_for_executorch,
328332
_load_for_executorch_from_buffer,
329333
)
330334
@@ -357,6 +361,99 @@ def forward(self, x, y):
357361
torch.testing.assert_close(actual.cpu(), eager, rtol=1e-3, atol=1e-3)
358362
capability = torch.cuda.get_device_capability(0)
359363
print(f"PASS: a CUDA-delegated model ran on sm_{capability[0]}{capability[1]} and matched eager")
364+
365+
366+
# The model above has no weights, so it needs nothing from the data file and runs whether or not
367+
# the delegate's external data path works. Anything a user would actually ship has weights, and
368+
# the CUDA delegate does not put them in the .pte: it writes them to a separate aoti_cuda_blob.ptd
369+
# that the caller has to supply at load time. Nothing here loaded one, so a wheel that shipped a
370+
# broken data path passed this check. Measured on sm_80: given the same model without the blob,
371+
# the delegate reports the weights as not found, and the run then dies inside forward with an
372+
# illegal memory access rather than returning wrong numbers.
373+
class Weighted(torch.nn.Module):
374+
def __init__(self):
375+
super().__init__()
376+
self.fc = torch.nn.Linear(8, 4)
377+
378+
def forward(self, x):
379+
return self.fc(x)
380+
381+
382+
weighted = Weighted().eval()
383+
weighted_example = (torch.randn(2, 8),)
384+
with torch.no_grad():
385+
weighted_eager = weighted(*weighted_example)
386+
387+
weighted_program = to_edge_transform_and_lower(
388+
torch.export.export(weighted, weighted_example), partitioner=[CudaPartitioner([])]
389+
).to_executorch()
390+
391+
assert b"CudaBackend" in weighted_program.buffer, (
392+
"the weighted program carries no CUDA delegate, so it would run on the portable kernels and "
393+
"never reach the external weights path this check exists for"
394+
)
395+
396+
# Only the weight matrix, not every parameter. Measured on sm_80, the delegate wrote the 128-byte
397+
# weight of this Linear into the data file and kept its 16-byte bias out, so the sum over
398+
# parameters() is not a bound on what the file has to hold.
399+
weight_bytes = weighted.fc.weight.numel() * weighted.fc.weight.element_size()
400+
401+
with tempfile.TemporaryDirectory() as work_dir:
402+
pte_path = os.path.join(work_dir, "weighted.pte")
403+
with open(pte_path, "wb") as handle:
404+
weighted_program.write_to_file(handle)
405+
weighted_program.write_tensor_data_to_file(work_dir)
406+
407+
blob_path = os.path.join(work_dir, "aoti_cuda_blob.ptd")
408+
assert os.path.exists(blob_path), (
409+
"the export wrote no aoti_cuda_blob.ptd, so the weights of this model went nowhere and "
410+
f"every runner script that passes that file by name would fail: {sorted(os.listdir(work_dir))}"
411+
)
412+
413+
# A data file is written even for a model with no weights, and it is not small: measured on
414+
# sm_80, the weightless model above produces a 256-byte container around an empty payload. So
415+
# the file's presence says nothing and its raw size says nothing either. What means something
416+
# is the amount by which this file exceeds that empty container, so write one and subtract it
417+
# instead of comparing against a constant, which would stop catching anything the next time
418+
# the container grows.
419+
empty_dir = os.path.join(work_dir, "weightless")
420+
os.mkdir(empty_dir)
421+
lowered.write_tensor_data_to_file(empty_dir)
422+
empty_size = os.path.getsize(os.path.join(empty_dir, "aoti_cuda_blob.ptd"))
423+
424+
blob_size = os.path.getsize(blob_path)
425+
assert blob_size - empty_size >= weight_bytes, (
426+
f"aoti_cuda_blob.ptd is {blob_size} bytes against {empty_size} bytes for a model with no "
427+
f"weights, a difference of {blob_size - empty_size}, which does not cover this model's "
428+
f"{weight_bytes}-byte weight matrix, so the weights were not written to it"
429+
)
430+
431+
with open(pte_path, "rb") as handle:
432+
pte_bytes = handle.read()
433+
with open(blob_path, "rb") as handle:
434+
blob_bytes = handle.read()
435+
436+
# The runner scripts pass this file by path, not as bytes, so exercise the
437+
# path-based loader here while the file still exists, or a regression in it
438+
# stays green.
439+
file_module = _load_for_executorch(pte_path, blob_path)
440+
file_actual = file_module.forward(list(weighted_example))[0]
441+
torch.testing.assert_close(
442+
file_actual.cpu(), weighted_eager, rtol=1e-3, atol=1e-3
443+
)
444+
445+
# Loading from bytes as well, because _load_for_executorch_from_buffer is what a
446+
# caller with the data already in memory uses, and it builds the data map by a
447+
# different route.
448+
module = _load_for_executorch_from_buffer(pte_bytes, blob_bytes)
449+
actual = module.forward(list(weighted_example))[0]
450+
451+
torch.testing.assert_close(actual.cpu(), weighted_eager, rtol=1e-3, atol=1e-3)
452+
print(
453+
f"PASS: a CUDA-delegated model with weights ran on sm_{capability[0]}{capability[1]} from a "
454+
f"{blob_size}-byte aoti_cuda_blob.ptd carrying {blob_size - empty_size} bytes of weights, "
455+
f"loaded both by path and from bytes, and matched eager"
456+
)
360457
"""
361458

362459

@@ -396,12 +493,17 @@ def _glibcxx_versions(library: Path) -> Set[str]:
396493

397494

398495
def test_a_model_runs_through_the_delegate() -> None:
399-
"""Export a model to the CUDA delegate and run it, comparing against eager.
496+
"""Export two models to the CUDA delegate and run them, comparing against eager.
400497
401498
Every other check in this file reads a shipped artifact. This one executes, because a wheel whose
402499
libraries are all present and correctly linked can still fail to compute, and nothing above would
403500
notice. The x86_64 rows land on a GPU runner, so this is the row where that can be proven.
404501
502+
The second model carries weights. The CUDA delegate keeps those out of the .pte and writes them
503+
to a separate aoti_cuda_blob.ptd, which the caller then has to supply, so a wheel can run the
504+
weightless model and still be unable to run anything real. That path is only covered by the
505+
second model.
506+
405507
Only the aarch64 rows may skip, and they say why, so a green result never stands for work that
406508
did not happen. On x86_64 the absent device or the mismatched torch build is itself the failure:
407509
that row is the one place execution is proven, and letting it report success for a check it did

0 commit comments

Comments
 (0)