Support Unbound Dynamic Tensor Memory
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 5k
- Forks
- 1.2k
- Avg merge
- 2d 10h
- Merged PRs (30d)
- 581
Description
🚀 The feature, motivation and pitch
Currently, all tensors must be statically memory planned, meaning they have a fixed upper bound and always allocate the maximum size, regardless of actual inputs. This is required to support certain constrained hardware, but is limiting for some use cases. I have also observed this behavior to be unintuitive to some engineers who are used to other frameworks. We have intended to support fully unbound tensors, but have not yet implemented it.
There are a number of use case which benefit from this:
- Models with wide input size ranges that don't want to pay the worst-case memory penalty unless necessary.
- Models where there is not a clear maximum input size. CV models which take an arbitrary input and resize in the model are a very common use case. They might want to be able to handle 4k images, for example, but don't want to statically allocate giant buffers unless necessary.
- Cases where users want it to "just work". They get frustrated with shape bounds and requirements. We can give them a way to opt-out and just deal with whatever inputs they send.
- Dynamic-sized KV caches?
- This can potentially be leveraged to implement overlapping of delegate allocations with ET-managed memory (see https://github.com/pytorch/executorch/issues/8764). This is a critical feature for XNNPACK to keep memory under control.
Some delegates, such as XNNPACK, already support unbounded tensors. CoreML supports it as an option, I believe. This helps bring the core runtime to parity. It also matches the experience with LI/TorchScript and some other ML frameworks where it "just works" with any size input.
As an additional consideration, doing this likely requires some form of runtime memory planning. Either that, or dynamic unbound tensor memory has to get freed when it goes out of scope. Otherwise, nothing will get overlapped and memory consumption will go through the roof. Imagine Llama 3 8B with unbound inputs size. There may be many ops on portable and not overlapping activations will be prohibitive.
Requirements
- Support dynamic unbound memory planning mode. Allow models to be lowered when exporting with unbound dynamic shape specification.
- Dynamic tensors should be able to be overlapped with other activation memory.
- Don't regress core runtime size for constrained targets. That might mean gating runtime memory planning behind a build-time configuration.
- Ideally, handle cases where tensor size is data dependent. Imagine instance segmentation models, where proposal count may vary.
Technical Solutions
Most of the technical complexity here is likely to lie in the runtime and memory planning. Export already supports unbound tensors, so no work should be needed there. Depending on the approach taken, we may need to serialize additional tensor lifetime data. The bulk of the work likely lies in implementing runtime memory planning and buffer reallocation + relocation.
Runtime Memory Planning
To make runtime memory planning work, we'll likely want to serialize tensor lifetime metadata in the .pte when runtime planning is enabled. We will need to allow for non-const buffers to be reallocated and relocated, which will require API changes. Ideally, this functions as a pure extension to the existing APIs, in order to preserve BC and leave statically-planned use cases alone. This needs be explored and scoped, and will likely be a sizable amount of work, but it should be do-able.
Alternative: Free tensors when out of scope
As an alternative to runtime memory planning, we can just allocate memory when tensors come into scope and free when it goes out of scope. The main concerns with this are:
- It won't be overlapped with statically planned tensors. Either we make everything dynamic (even static tensors), or we risk a huge memory increase. This alone may rule out this option.
- It carries a performance penalty, though I don't have a feel for how much.
- It can't be used to overlap delegate memory with core (unless everything is dynamic). See https://github.com/pytorch/executorch/issues/8764.
Because of these options, I strongly prefer runtime memory planning.
Alternatives
No response
Additional context
No response
RFC (Optional)
No response
cc @JacobSzwejbka @mergennachin @byjlw
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
No specific files or tests are named. Start by reading the existing export path, runtime memory-planning implementation, .pte serialization, and tensor buffer APIs; scope how unbound tensors can be relocated and overlapped without affecting statically planned targets. Done means unbound export works, dynamic activations overlap safely, constrained runtime size does not regress, and data-dependent tensor sizes are handled.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- embedded-iot, machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100