Improve memory usage in EXIR emitter

[iseeyuan]

• Parameters in json file, that the json str is 5x larger than the file size. One option is to have parameters in segments in [internal T165017481]
• Redundancy of different artifacts. Before the emitter, the memory size is > 2x larger than the model size.

[lucylq]

Flatbuffer Serialization for Models with Large Constants

Objective

To export models with large constant data (weights, biases) in ExecuTorch.

The Llama model can contain 7b+ parameters. Currently, this causes an OOM during serialization. [Internal repro T163016365].

Background

serialize_pte_binary performs the following:

• Extract segments from the program
• Convert Program to JSON
• Convert JSON to flatbuffer
• Compile the .pte file together; insert header, attach segments

Currently, constants are stored in the constant_buffer field of the program. When we convert the program to JSON, these tensors are converted to string and occupy much more space. This causes an OOM error when constant data is large.

There are two ways we can potentially improve upon the export/serialization design to resolve this problem.

Avoid converting to JSON

Previously, the flatbuffer Python API was not well-supported and was missing features (eg. union), which is why the current implementation converts to JSON. Since then, support has improved and we can potentially use the flatbuffer Python API to avoid the conversion to JSON. This resolves the OOM error as tensors are not converted to JSON.

Move constant data into segments

The constant data inside the Program is serialized tensor data that doesn’t need to be encoded inside the flatbuffer. We can move the constant data into segments (for example, with delegate data blobs). Segments do not go through JSON conversion and flatbuffer serialization; they are attached at the end of the .pte file. This resolves the OOM error as tensors bypass the conversion to JSON.

Evaluation

We propose to move constant data into segments. This places constant data with the delegate segments, both of which contain serialized data.

Flatbuffer has a size limit of 2GB, so purely updating to use the flatbuffer Python API may not resolve the issue for models with really large tensor data.

In the future, we can consider moving to the flatbuffer Python API to streamline the Program→JSON→flatbuffer into Program→flatbuffer. This has its own considerations; it may make schema changes more difficult as we’d have to edit serialization code (today, we update the schema and JSON will be generated in that form), and we’d need to ensure that serializing directly with flatbuffer Python provides the same results (though this should be guaranteed by flatbuffer).

Goals

Store weights/biases/other constants that are currently stored in constant_buffer, in a separate segment, like how delegate segments are stored today.
Ensure that the runtime can parse the new flatbuffer schema
Ensure the runtime supports the old flatbuffer schema (not BC breaking)

Design

We propose to move constant tensor data currently stored inside the flatbuffer, into a new segment that is appended to the flatbuffer data. At a high level, this involves:

• AOT: introduce the capability to serialize constant data into a separate segment.
• Runtime: add capability to load both new and old serialization formats.

When constant data is stored in a separate segment:

• The existing field constant_buffer* will be empty and unused
• The new fields constant_segment.segment_index* and constant_segments.offsets* will be non-empty, unless there is no constant data.
• *See diagram below

And vice-versa when constant data is inlined in the flatbuffer.

Tensor.constant_buffer_idx() values, used to reference the tensor information, will point into either constant_buffer or constant_segment depending on whichever one is present in the .pte file being loaded.

AOT Changes

• Optionally move tensor data from constant_buffer to segment
• Store segment index and tensor offsets in newly added constant_segment

program.fbs defines the following schema objects:

table Program {
  version:uint;
  execution_plan:[ExecutionPlan];
  constant_buffer:[Buffer];
  backend_delegate_data:[BackendDelegateInlineData];
  segments:[DataSegment];
}
table DataSegment {
  offset: uint64;
  size: uint64;
}
table Buffer {
  storage:[ubyte] (force_align: 16);  // @executorch-tensor-alignment
}

We introduce a new field:

table Program {
  //...
  constant_segment: SubsegmentOffsets;
}

table SubsegmentOffsets {
  segment_index: uint64;
  offsets: [uint64];
}

Move constant_buffer to segment

Each Buffer in constant_buffer contains a separate tensor. To move them into a single segment:
Align each tensor to constant_tensor_alignment bytes, by default 16.

Store the segment offset and size in segments list.

There will be one segment for all methods in the program. This enables data sharing between methods, and mirrors the current state where all the constants are together.
Index constant segment using constant_metadata
We store the segment index and tensor offsets in the newly introduced constant_segment field. The offset of each tensor within the segment is aligned to 16 bytes.

Runtime Changes

To ensure BC, the runtime will have two branches to either load the new or old schema. If constant_buffer and constant_segment are both empty, then the model probably doesn’t have any constant data. Otherwise, we load assuming the old schema if constant_buffer contains items, or the new schema if constant_segment contains items. If both contain items, we throw an error.

Both branches will exist until we decide (if we decide) to deprecate constant_buffer.

Program Load

The constant segment should be loaded during program load time, so tensors can be freely accessed by each method (avoid reloading for each method). Note that delegate segments are loaded at method load time.

Method Load

To access the tensor at index n:

• Retrieve segment_index from constant_segment
• Retrieve tensor_offset from constant_segment[n]
  Tensor data is located at `segments[segment_index] + tensor_offset

Note; only require the offset (not the size), as the tensor is pre-allocated with the correct size.

BC/FC, Versioning, Compatibility

✅BC: the new runtime can load models exported using the old schema.

• Changes proposed are BC safe, as the runtime will branch and decide which schema version to load. No schema version bump is required now.

❌FC: the old runtime can not load models exported using the new schema.

• Changes are FC breaking, as the old runtime won’t see the new constant_metadata field and won’t be able to resolve const indices.

[MichaelTylerArm]

This looks like a promising solution to the OOM issues when converting large models, although I have personally experienced other issues with the JSON conversion that would make it desirable to avoid if possible. In particular, on my system flatc seems unable to read or write the temporary files used by ExecuTorch, so I had to hack it to operate in the current working directory instead. I don't know if this is just some configuration issue with my OS (Ubuntu), but I imagine I am not alone in this and a conversion pipeline that doesn't need to work with temporary files and external commands would be preferable.