Uh oh!
There was an error while loading. Please reload this page.
Add TensorRT weight streaming support to the ExecuTorch delegate - #4336
Conversation
983583d to
053e902CompareGood call, and I agree the budget should be a runtime setting. Here is what the PR does today, and a change that gives you exactly what you are asking for. What happens today
So the only thing fixed at export is the explicit number, and you are right that a fixed byte count is really a per-deployment choice. Proposed change (a real load-time API)
Why keep the export-time value too ExecuTorch's Python and Android load paths do not expose backend options yet (only C++ Changing it after load (the "or later" part) This is doable as a follow-up, not in this PR. TensorRT requires destroying and recreating the execution context to change the budget, and ExecuTorch's post-load One question so I build the right thing Will these large models be served mainly from the C++ runtime or from Python? If C++, the load-time option covers it and we can treat the baked value as just a default. If Python, we need to keep the baked value until ExecuTorch exposes backend options to Python, which I am happy to help add upstream. |
c354c29 to
a996b54Compare
cehongwang
left a comment
There was a problem hiding this comment.
Overall OK. Some minor comments
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
3734e22 to
91ae41aCompareshoumikhin
commented
Aug 18, 2026
Both of these are addressed now. Empty runtime option. An empty value meant "unset", and the code fell through to the Budget with a graph break. Good question to push on, and the answer is that there is no While testing that, the type check for the budget turned out to run after the model type check, |
shoumikhin
commented
Aug 19, 2026
It is settable when the method loads: a valid non-empty On "or later": there is no post-load setter here, because TensorRT needs active execution |
Uh oh!
There was an error while loading. Please reload this page.
cehongwang
left a comment
There was a problem hiding this comment.
Just one warning issue. And it is good to go
d1e9fde to
82c9e2dCompare… delegate TensorRT can build an engine that keeps only part of its weights in GPU memory and streams the rest from host memory while the engine runs. That is what lets a model whose weights do not fit on the device still run there. Up to now the ExecuTorch delegate had no way to say how much GPU memory those weights may use, so a streaming-capable engine always got TensorRT's own automatic budget. This adds an explicit budget you can set at export time: torch_tensorrt.save( program, "model.pte", output_format="executorch", weight_streaming_budget_per_engine=8 * 1024**3, ) torch_tensorrt.executorch.export( program, weight_streaming_budget_per_engine=8 * 1024**3, ) The value is bytes. It is baked into the delegate as a compile spec and read back by the C++ backend when the method loads. The same key is also accepted as a load-time backend option, which wins over the baked value when both are present. Leaving it unset, the default, keeps TensorRT's automatic budget, sized against free memory on the device the engine actually loads on. The budget applies to each engine on its own, not as a total for the program, which is why the name says "per engine". A program that splits into N engines can hold up to N times the value resident, because every delegate is initialized when its method loads and they stay resident together. Export logs a warning when a program has more than one engine and a budget is set. An engine that was not built with enable_weight_streaming cannot stream, so the budget is ignored there and the delegate logs why. Passing the raw CompileSpec by hand is rejected, because that route skips the validation and can silently disagree with the argument. Tested: new CPU-only Python tests covering budget validation, the compile spec reaching the TensorRT partitioner for single-method and multi-method exports, the rejection of a hand-written spec, the multi-engine warning, and the save() argument guards. New C++ tests covering the backend option parsing and precedence. Both were checked to fail when the corresponding code is removed.
82c9e2d to
bf28635CompareUh oh!
There was an error while loading. Please reload this page.
A multi-layer model lowered to the ExecuTorch TensorRT delegate becomes N
separate single-layer engines, each with its own execution context. By default
every context allocates its own activation scratch (`getDeviceMemorySizeV2`
bytes) and holds it for as long as the context lives, so device memory scales
with the layer count and multi-layer models OOM at runtime. The scratch of one
engine need not sit alongside the scratch of the next, because the delegates of
a Method are submitted one at a time: `Method::execute()` advances `step_state_`
through the instruction stream one instruction at a time on the calling thread,
and a `DelegateCall` is one instruction. Their enqueues can still overlap on the
device, but that is orderable, whereas N resident copies are not.
Submission order is not stream order, though. That two consecutive delegates
land on the same stream is not a property of `Method`; it holds only because
both read the same thread-local caller stream, and a caller that runs two
Methods under two different `CallerStreamGuard` streams breaks it. So the pool
orders the handoff itself rather than relying on the stream.
Add an opt-in shared pool that backs all contexts on a device from one buffer:
- the `use_shared_activation_scratch` runtime backend option enables it. It is a
boolean, defaults to false, and is delivered with
`executorch::runtime::set_option("TensorRTBackend", options.view())`. With it
unset each context owns its private `kSTATIC` scratch, the delegate emits no
extra log output, and the only work it adds is one relaxed atomic load per
engine init and a bool test or two per `execute()`;
- when enabled, create each execution context with `kUSER_MANAGED` so it
allocates no scratch of its own (`initialize_engine_io`);
- in `execute()`, once the input shapes are bound, query the exact requirement
with `updateDeviceMemorySizeForShapes()`, grow a per-device pool to it, and
point the context at the current buffer via `setDeviceMemoryV2`.
The pool grows monotonically to the largest engine's need and syncs the device
before freeing a replaced buffer. N per-layer scratch copies collapse to one
(the (N-1)x duplication is reclaimed). Measured with TensorRT 11.2.1.2 and CUDA
13 on one 80GB NVIDIA PG509-210, in a CMake reference runner that also loads the
ExecuTorch CUDA/AOTI backend, reading `cudaMemGetInfo` after a `cudaFree(0)`
baseline, on a deterministic non-uniform fp32 input:
- four execution contexts of one engine holding two fp32 8-head attention blocks
over `[1,2048,512]` (285,212,672 B of scratch) go from 1188MB to 372MB,
3 x 272MB reclaimed;
- a single Method holding six one-block engines of the same shape, interleaved
with six CUDA delegates (281,018,368 B each), goes from 1656MB to 316MB,
5 x 268MB reclaimed.
Outputs are identical between the two modes in both cases.
Two consequences a caller feels once the option is on. The pool is never freed,
so a device keeps the largest scratch it was ever asked for until the process
exits, where per-context `kSTATIC` scratch is released with its context. And a
growth allocates the new buffer before releasing the old one, so both are
resident for that moment -- that ordering is what leaves the existing buffer
usable when an allocation fails.
An execution context's allocation strategy is fixed when the context is created,
so each engine captures the setting in effect at its own init and keeps it. A
later `set_option` decides what the engines loaded after it are built with and
changes nothing about the ones already running, so a `kSTATIC` context and a
`kUSER_MANAGED` context coexist in one process.
Why opt-in, not default-on: one buffer serves every context on a device, and a
context holds its scratch for the whole enqueue -- which under a
`CallerStreamGuard` can still be in flight when `execute()` returns -- so two
enqueues must never hold it at once. The pool records each enqueue on a
per-device event and makes the next one wait on it. An event, not the previous
stream: synchronizing on a destroyed stream handle crashes rather than returning
an error; CUDA recycles handle values, so two distinct streams can compare
equal; and the NULL stream is a legal caller stream that no stream-handle
sentinel can tell from "no previous user". Waiting from the stream that recorded
the event is already satisfied, so the single-stream case pays a host call and
no device stall. Not covered: concurrent same-device `execute()` on several
threads, because the pool mutex is released before either enqueue is submitted.
A default-on version needs the scratch keyed per stream instead of one buffer
per device.
This is orthogonal to weight streaming (pytorch#4336), which targets engine *weight*
memory rather than activation scratch, and to export-time OOM.
The grow/reuse/per-device policy and the handoff rule are factored into a
header-only helper (`SharedScratchPool.h`) so they are unit-tested without a
device (fake allocator, fake event factory). The CUDA path supplies the three
callables it takes -- `cudaMalloc`, `cudaFree` and `cudaEventCreateWithFlags`;
the `cudaStreamWaitEvent` and `cudaEventRecord` half of the handoff is the
caller's, issued in response to what the helper returns. The helper needs the
`cudaEvent_t` typedef, which is a compile-time dependency and not a runtime one;
the repo had no headers-only CUDA target, so `third_party/cuda/BUILD` gains one
rather than the helper depending on `cudart` and putting `libcudart` in the
DT_NEEDED of a test that makes no CUDA call. `set_option` itself is not
unit-tested, because no target in `tests/cpp/executorch/` links the backend.
Three behaviours live only there: skipping a key this backend does not read,
storing a valid boolean, and rejecting a non-boolean with
`Error::InvalidArgument` instead of dropping it silently. The store is exercised
by the measurement above, which reaches the pool through `set_option`; the key
skip and the wrong-type rejection are covered nowhere, as `CudaBackend`'s
equivalents also are.
Summary
Add TensorRT weight streaming support to the Torch-TensorRT ExecuTorch delegate, so a model whose weights do not all fit in GPU memory can run when exported to an ExecuTorch program. Ref #4334.
Torch-TensorRT already builds a weight streamable engine when you compile with
enable_weight_streaming=True, but the ExecuTorch delegate never set a budget on the engine at load time, so large models could not stream. This change sets the budget in the delegateinit(), after the engine is deserialized and before the execution context is created, which is the same pattern the other Torch-TensorRT runtimes already use.How it works
By default the delegate applies TensorRT's automatic budget, computed at load time from the free memory on the actual GPU, gated on
getStreamableWeightsSize() > 0. So an engine built withenable_weight_streaming=Trueruns out of the box and adapts to the deploy device. Nothing is baked into the.ptefor this default case.An explicit budget is a non-negative number of bytes and can be set two ways, in order of precedence:
weight_streaming_budget, passed by the caller viaModule::load(LoadBackendOptionsMap)and read ininit()withBackendInitContext::get_runtime_spec. This lets a deployment size the budget for its own GPU without re-exporting. It is the same load-time pattern CoreML and XNNPACK use..pteviatorch_tensorrt.save(output_format="executorch", weight_streaming_budget_per_engine=N). The Python keyword is named_per_enginebecauseweight_streaming_budgetalready means a program-wide total onMutableTorchTensorRTModule; the key on the wire is unchanged. This is used when no load-time option is given, and it is the only channel for loaders that cannot pass backend options yet (the ExecuTorch Python and Android runtimes).Resolution order in the delegate is: load-time option, then the baked value, then automatic. The value is a decimal string on the wire because ExecuTorch's typed integer option is only 32 bit and a byte budget can exceed 2 GB.
Changes
TensorRTBackend::initapplies the budget viasetWeightStreamingBudgetV2before creating the execution context, gated ongetStreamableWeightsSize() > 0. It resolves the budget as load-time runtime spec, then baked compile spec, then automatic.WeightStreamingBudgetparser (cpp/includeandcpp/src), unit tested without a GPU. It usesstd::from_charsand accepts only a non-negative decimal integer.torch_tensorrt.executorch.export(..., weight_streaming_budget_per_engine=...)writes the export-time default into the TensorRT partitioner's compile specs, for every method of a multi-method export.torch_tensorrt.save(output_format="executorch")forwards the same argument. Validation lives next to the compile spec key intorch_tensorrt/executorch/partitioner.py. Passing the budget throughcompile_specsby hand is rejected in favor of the keyword argument, because that route skips the validation.Requirements
The load-time override uses ExecuTorch's
BackendInitContext::get_runtime_specandLoadBackendOptionsMap. The export-time default and the automatic budget work without it.Backward compatibility
.pte..pteformat or the engine blob.Edge cases
Error::InvalidProgram).weight_streaming_budget_per_engineasNonefor multi-engine models.Status and validation
reaching the TensorRT partitioner for single-method and multi-method exports, the rejection of a
hand-written spec, the multi-engine warning, and the
save()argument guards. Each was checked tofail when the code it covers is removed.
tests/py/dynamo/executorchdirectory on this change and on the base branchgives the same set of failures, so the change adds no new ones. Those shared failures are from the
test machine having no GPU and no compiled Torch-TensorRT runtime, not from the code.
by CI rather than by hand here.
Follow-ups
LoadBackendOptionsMapin ExecuTorch's Python (and Android) runtime bindings so non-C++ loaders can also set the budget at load. Until then the export-time default covers them.set_optionAPI. This needs the execution context to be destroyed and recreated, likeTRTEngine::set_device_memory_budget, so it is deferred.init()with a test: runtime-spec versus compile-spec precedence, the clamp togetStreamableWeightsSize(), and the automatic-budget fallback. The current tests cover the byte-string parser and the export-time plumbing only. This needs an engine actually built withenable_weight_streaming=True, so it belongs in the GPU delegate lane rather than here.Behavior change worth calling out
save(output_format="executorch", ...)now raisesTypeErroron an unrecognised keyword argument instead of ignoring it. This turns a silent typo into a hard failure, which is the point, but it is a behavior change for anyone who was passing an unused kwarg. It is scoped tooutput_format="executorch", so the other output formats are unaffected.