[aoti-et] Store weights outside of .so

[larryliu0820]

This pull request introduces significant changes to how model constants (weights) are packaged and loaded in the CUDA backend, improving modularity and flexibility. It also updates the runner APIs to support more flexible loading modes and bumps PyTorch commit pins and nightly version references.

Packaging and Loading Model Constants:

• Model constants are now separated from the .so file and stored as a binary blob on disk, rather than being packaged directly into the shared object. The preprocessing logic in cuda_backend.py is updated to handle the new file outputs and manage cleanup. [1] [2]
• The C++ CUDA backend (cuda_backend.cpp) now loads the new weights blob from the named data map and feeds it into the model container using a newly added API function. The buffer is freed immediately after use for better resource management. [1] [2]

API and Infrastructure Updates:

• A new function pointer type, AOTInductorModelUpdateConstantsFromBlobFunc, is added to the delegate handle structure in aoti_delegate_handle.h to support updating model constants from a binary blob. [1] [2]
• The CUDA backend now loads this new symbol from the shared object at runtime.

Runner API Improvements:

• The multimodal runner API is updated to accept a Module::LoadMode parameter, allowing for more flexible loading options such as memory mapping. This change is propagated through helper functions and their headers. [1] [2] [3] [4]

Dependency Updates:

• The PyTorch commit pin is updated in .ci/docker/ci_commit_pins/pytorch.txt and the nightly version is bumped in torch_pin.py for compatibility with the new packaging logic. [1] [2]

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/15180

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 77 Pending, 1 Unrelated Failure

As of commit 989def3 with merge base 4cff294 ():

NEW FAILURE - The following job has failed:

• pull / unittest-nxp-neutron / linux-job (gh)
  RuntimeError: Command docker exec -t 359c50fad57cfdaa72c2259c028402f905b758688e93ba3f72a409936204b680 /exec failed with exit code 1

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• docker-builds / docker-build (linux.4xlarge, executorch-ubuntu-22.04-gcc9) (gh) (detected as infra flaky with no log or failing log classifier)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[mergennachin]

check if existence/non-existence of blob_path matches the options

[mergennachin]

Only check for blob_path is None if package_constant_on_disk_format option is set to "binary_blob"

[swolchok]

the weights are mmapped, so this isn't halving the maximum amount of weights we can handle, right? even so, seems unfortunate that we have to copy and therefore can't keep them simply mmapped though; peak CPU memory now needs to hold them, right?

[larryliu0820]

this isn't halving the maximum amount of weights we can handle, right?

Good point, let me test this on my RTX 5080.

seems unfortunate that we have to copy and therefore can't keep them simply mmapped though

Yeah would be good if aoti can just take it without copying.

[yushangdi]

Yeah would be good if aoti can just take it without copying.

I might be missing something here, but I assume you're mmaping into the CPU memory right? AOTI copies it into the CUDA memory, and since we're running this on CUDA, we have to copy it to CUDA some time.

[larryliu0820]

Is there a mmap equivalent on CUDA? If so, on ET side we can create a dataloader to directly load into CUDA memory.

[yushangdi]

I guess you can use GPUDirect Storage.