Memory Profiling in DevTools

[GregoryComer]

🚀 The feature, motivation and pitch

Measuring inference and load-time memory consumption is a common task when looking to use ExecuTorch. It's useful to determine if a model can actually run on consumer hardware and to check parity against other inference solutions. Users can profile their entire app, but it takes a bit of work to figure out what is owned by ET, and it doesn't provide any detailed breakdown of where the memory usage is coming from.

We have performance profiling in the ExecuTorch devtools. It would be good to also have memory profiling support. What exactly is captured at runtime needs some discussion, but as a user, I'd ideally want to know the following:

• What is the peak memory used by the framework (including delegates) during inference?
  • How much is owned by ET?
  • How much is owned by delegates?

Some nice to haves:

• How much (peak) memory did we use during model load?
• Can we take RSS delta measurements on supported systems?

There is some overlap between this and memory planning visualization. In my opinion, memory profiling tooling fills an important role that static memory planning visualization can't when it comes to delegates that dynamically allocate memory, and for tracking other dynamic memory allocations (such as from kernels or dynamic unbound tensors).

Alternatives

No response

Additional context

No response

RFC (Optional)

No response

cc @Jack-Khuu

[iseeyuan]

I remember there were discussions on this topic. In general, ExecuTorch runtime is static.

other dynamic memory allocations (such as from kernels

There's temp buffer passed to kernels via context. It's decided AOT.

or dynamic unbound tensors).

That's true. But I'm wondering if we have unbound tensors in any use cases.

delegates that dynamically allocate memory

I'm wondering what are the situations of dynamic memory allocation in delegates? For NPUs it would be static in most cases, or did I miss anything for either QNN (cc @cccclai ) or coreML (cc @metascroy ). For xnnpack, there's prepacking that may take extra memory, but would it be estimateable AOT? There's "workspace" that may take more memory. @digantdesai and @mcr229 , do you know if it's possible to estimate it AOT?

[tarun292]

There are 2 aspects to this:

1. Allocations that are done through the allocators we provide
2. Allocations that are done outside of these allocators (via malloc etc.)

For 1. we already have support in etdump to log this, but we just haven't piped this in the runtime. This'll give us peak memory consumed from the allocators for activation buffers, kernel buffers from temp allocators, memory for loading methods etc.

2. will be more relevant for delegates etc. and that's harder to do as these are done via various system level allocators and these can vary a lot from system to system. It's quite hard for us as a framework to tap into this but rather it'd be dependent on the users to use their system level memory profiling tools to track this peak memory consumption.

[GregoryComer]

For XNNPACK specifically, all of it's allocations are dynamic. CoreML presumably works similar. XNN can take in a custom allocator, which we can use to track allocations. I have a hacky version of this I've used extensively. For XNNPACK-delegated models, the XNNPACK library typically owns most or all of the activation memory. Other delegates may manage their own memory, which we should look into. In general, right now, most of the runtime memory on mobile delegates isn't owned by ET, but rather by the delegate. It would be interesting to here more on if CoreML, Qualcomm, and other delegates can potentially tap into this.

We could build an XNNPACK-delegate specific tool here, but it makes a lot more sense to me to try to generalize this to other delegates. I believe this was already in the pipeline for devtools - I just created a GH issue to track it publicly. IMO it makes sense to offer delegates a way to opt-into this, where we can provide a custom allocator to the delegate to use, which will integrate with profiling. Using system-level profiling has limitations, especially when you want more detailed breakdowns on why allocations are made. It needs more discussion, but we could potentially give the ability to flag an allocation as weights / activations / etc.

I've actually spent much more time fighting memory bloat than performance for internal model launches, which is part of why I think we might want this. Users have also asked for this, as well, and have been surprised when I told them that we don't have tooling for it.

[Gasoonjia]

For static memory usage (the memory usage planed during memory planing) we already have a tool for profiling: https://pytorch.org/executorch/stable/memory-planning-inspection.html

For memory allocations outside memory planning, I think we can provide an interface which can be extended by users for their custom memory tracing and also be used in the etdump/etrecord for further inspection.