RFC: PTE Size Inspector Design

[Olivia-liu]

🚀 The feature, motivation and pitch

Problem

Currently, users who export models to ExecuTorch have no tool to inspect what contributes to the size of the resulting .pte file. This is a concern because the file must fit within the available memory on devices, which are often very limited.

Goal

Users can understand what contributes to the overall size of a PTE size using a commandline tool or python script/notebook.

RFC

Design

Overview

Have 3 entry points to inspect .pte file size:

• Before having the .pte:
  • Entry Point 1: Python util function size_distribution(exec_prog). Use at the end of the export script.
• When already have the .pte:
  • Entry Point 2: Command Line tool, pteinspect, to Inspect the .pte file. Use in a terminal.
  • Entry Point 3: Python util function size_distribution_from_pte(pte_file). Use in a python script or in a python notebook.

In order to get detailed sizing information in the delegate blobs, allow delegates to implement hooks to decrypt the delegate blob. It’s optional for the delegate authors to implement. See the comments below for more discussion on this.

Details

Class SizeDistribution and size_distribution, size_distribution_from_pte util functions

SizeDistribution is a recursive data class designed to hold size distribution information. It also comes with SizeDistribution.to_dataframe() to get size distribution details in the format of a pandas dataframe.

size_distribution are convenient util functions to get a SizeDistribution instance from an ExecutorchProgramManager instance or a .pte file.

User Interface

# User code: in export.py or in a notebook, call size_distribution() after to_executorch() is called in the export process
from executorch.devtools.ptetools import size_distribution

...
exec_prog = edge_program.to_executorch()
size_dists = size_distribution(exec_prog)
print(size_dists)
# If want to use dataframe
df = size_dists.to_dataframe()

# User code: in a notebook or a python script file, call size_distribution() after user already has a .pte file
from executorch.devtools.ptetools import size_distribution_from_pte

file_path = "/path/to/your/.pte"
size_dists = size_distribution_from_pte(file_path)
# Rest is the same as the example usages above

Example output of print(size_dists)

Program Flatbuffer: 51.23 KB
Constant Tensors: 112 B
	conv_0_weight: 64 B
	conv_1_weight: 48 B
Delegate Blobs: 13.90 MB
	XnnpackBackend_0: 8.85 MB
	XnnpackBackend_1: 5.05 MB

Print in human-readable scale units:

• If size >= 1 GB, use GB
• If size >= 1 MB, use MB
• If size >= 1 KB, use KB
• If size < 1 KB, use B

Example df when printed out

Name	Size (bytes)	Level
Total Size	13951248	0
Program Flatbuffer	51232	1
Constant Tensors	112	1
conv_0_weight	64	2
conv_1_weight	48	2
Delegate Blobs	13900016	1
XnnpackBackend_0	8845856	2
XnnpackBackend_1	5054160	2

Implementation

# New file: executorch/devtools/ptetools.py


class SizeDistribution:
	"""Sizes represented in a recursive structure"""
def __init__(
self, 
name: str, 
size: int, 
components:  Optional[List['Size_Distribution']] = None
):
       	self.name = name
       	self.size = size
       	self.components = components 

	def __str__(self, level=0):
       """String representation for displaying the hierarchy with indentation."""
       indent = "  " * level
       result = f"{indent}{self.name}: {self.size} bytes\n"
       for component in self.components:
           result += component.__str__(level + 1)
       return result

def to_dataframe(self):
	""" Format the class into dataframe """


def size_distribution(exec_prog: ExecutorchProgramManager) -> SizeDistribution:
	"""
	Args:
	    exec_program: ExecuTorch program
Returns: 
    Hierarchical size distribution of different components of exec_program 
	"""
	
def size_distribution_from_pte(file_path: str) -> SizeDistribution:
	"""
	Args:
	    File path of a .pte file
	Returns:
	    Hierarchical size distribution of different components of the .pte 
	"""

Command Line tool, pteinspect

This is useful for users who don’t necessarily export the model themselves, and have a .pte and want to understand the size of it. Users can also call this from a bash script to do pte file analysis.

Example user flow:

$ pteinspect [options] pte_file

where,

• [options]: Various command-line options that determine the output and level of detail displayed by pteinspect.
• pte_file: The PTE file to be inspected.

Option	Description
-l	List the top level components of the .pte file and the size of each of them
-e component_name	List all the components inside component_name and the size of each of them
-h	Display information in the headers
--help	Provides a help message listing all available options

Alternatives Considered

Considered having an interactive commandline tool, but decided to move away from it because a commondline with arguments is more scriptable, and the style also matches more with ELF tools, which is widely used in the industry.

Also considered combining different pte tools (file inspection, file modification, etc.) into one tool. Decided to have separate tools for different features to match with ELF tools style, and also give users confidence that they wouldn’t accidentally modify the file when they only want to inspect it.

Release Plan

Milestone 1 (1 week): Define Python class and write Python APIs

Milestone 2 (1 week): Write the commandline tool

[Olivia-liu]

Delegate Blob Hook Discussion

The .pte file can have multiple segments and in them there could be multiple delegate blob segments, each represents a graph description (for the delegated subgraph) and constant tensor data. Users would be interested in knowing what’s being included in the delegate blobs. We should add hooks that delegate authors can potentially implement in order to return such information.

While exporting is in progress:
The to_backend pass is where the delegate blobs are created by the partitioners. I propose adding a step in the partitioning process to construct a dict in which the delegate blob details are written to. And while exporting, we not only produce a .pte, but also a metadata.json of each delegate blob.

After exporting:
If the .pte has already been created, and during to_backend no extra step was taken to get metadata.json for delegate blobs, users should still have a chance to inspect the delegate blobs. I propose for each partitioner to have a new .pte parsing function with which a metadata.json can be reverse-engineered from a delegate blob that this partitioner had created.

[Olivia-liu]

@digantdesai @dbort @larryliu0820 Would you be able to give feedback on the "Delegate Blob Hook Discussion" comment above? And would you be able to sign off the design (by commenting)? Thank you!

[digantdesai]

This is awesome @Olivia-liu. Thanks for putting this here.

For the top level APIs, a couple of nit comments,

1. The python APIs are more aptly named size_distribution[_from_pte] but the command line tool is named pteinspect, that seems a bit vague to me especially if we aligned on keeping it a narrow focused size tool. Perhaps, we may be consider pteinspect size -l as a subcommand instead of pteinspect -l?

2. Please include an example output of the pteinspect -l tool to get some feedback about output format (for parsing and so on).

Otherwise, the proposed top-level APIs/tools look good to me.

For the delegate hooks design discussion, I feel we may get get more meaningful feedback if we provide some more details (even if not fully solidified). For example, for AoT, what the delegate.preprocess() generated metadata dict might look like, or the metadata.json for that matter. Similarly for runtime, what is the output of the delegate authored blob parser might look like.

Besides that, I think the delegate hooks, at a high level, as proposed here looks good to me. It is opt-in, and when opted for, gives delegate implementor a choice for implementation.

[dbort]

Thanks Olivia, this looks useful!

Comments on the details of the proposal:

• I suggest renaming the size_distribution functions to something like size_distribution_from_buffer and size_distributions_from_file, and make the _from_buffer take a Union[bytes,bytearray,Cord] that contains the PTE data. Users would then call it like size_distribution_from_buffer(program.buffer), where the buffer contains the PTE data. This would make the API abstraction more consistent: it always operates on PTE data, whether it's in a file or in a buffer. And it also reduces the dependencies of ptetools, since it wouldn't need to depend on ExecutorchProgram.
• For human-readable sizes:
  • Kilobytes should be small k: kB, like the k in km
  • Are these in multiples of 1000 or 1024? If 1000 then the abbreviations are ok. If 1024, then they should be kiB, MiB, GiB etc. See https://simple.wikipedia.org/wiki/Mebibyte vs. https://simple.wikipedia.org/wiki/Megabyte. I suggest using 1024 because that's what most unix tools do: e.g., see the --human-readable flag for the df utility. And using the MiB form makes it more clear that 1024 is the divisor.
• I agree with Digant's recommendation to put this functionality under a size subcommand to make room for other non-size-related pteinspect features in the future.
• Commandline options:
  • It is generally better if all options have long names, and optionally have short names.
  • Please describe the default behavior when passing no options. Are they all enabled? All disabled? It's nice if the default behavior gives something useful, like a summary.

For the commandline output:

• If the intent is that a bash script could parse the output, the current output is tough to parse. It requires contextual knowledge: what was the most recent section string? "Constant Tensors:" or "Delegate Blobs:"?
  • One solution for this is to avoid section lines, and to make each line self-describing E.g. CT conv_0_weight 64 (CT for constant tensor) or DB XnnpackBackend_0 9279897 (DB for delegate blob).
    • Note that in this case I removed the MB part, because that's also hard to parse; raw numbers are easier when parsing, but they're definitely useful when humans are reading the output.
    • Another benefit to a format like this is that users could pipe the output to sort -k3 to sort the entries by their sizes.
  • Another solution would be an option that prints the data in a more parseable format like JSON or CSV.

Some higher-level user flow questions:

• What do you expect is the next step for a user who uses this tool to find an unusually large tensor/blob/etc.? How do they use the output of this tool to e.g. update a model to reduce memory usage? Is there a clear path from an entry in the output to a line of python in the nn.Module?
• A lot of models have thousands and thousands of tensors and graph nodes. Right now, it seems like the two options are: sums for "flabuffer", "constant tensors", and "delegate blobs", or a list of every tensor/blob in the file. How would a commandline user narrow down the lines they need to look at? Do you expect to add filters, sorting, or anything like that?

[larryliu0820]

A lot of models have thousands and thousands of tensors and graph nodes.

I agree with @dbort here that this is a valid concern. Following this train of thought:

Do you expect to add filters, sorting, or anything like that?

I would suggest some more features/options for the command line tool:

1. By default sort the sizes in descending order
2. Only show top 5 (or 10?) by default and add an option to expand and show all values
3. Add an option to filter by a threshold - anything larger than X KiB

[digantdesai]

I would suggest some more features/options for the command line tool:

I guess if we format output correctly, like --human-readable then we can easily implement what you are suggesting using standard unix-way, and we can open up more options too, for instance,

(1) pteinspect size --human-readable x.pte | sort -hr # for sorting
(2) pteinspect size x.pte | head -n 10 # for first 10
(3) This can be an argument TBH :)