10ba1e270c
🛠️ __Changes in this pull request:__ This pull request introduces two significant changes to the project: - Changing on device training checkpoint format: The current implementation stores the on device training checkpoint as a sequence of tensors in multiple files inside a checkpoint folder, which can be inefficient in terms of storage and performance. In this PR, I have modified the checkpoint format to utilize the flatbuffer table to save the checkpoint to a single file, providing a more compact and efficient representation. The changes around this are twofold: - Add the checkpoint flatbuffer schema that will generate the necessary checkpoint source files. - Update the checkpoint saving and loading functionality to use the new format. - Adding support for onnxruntime minimal build: To support scenarios where binary size is a constraint, I made changes to ensure that the training build can work well with the minimal build. 🔍 __Open Issues:__ - In order to extract the optimizer type, the existing implementation re-loaded the onnx optimizer model and parsed it. This is no longer possible, since the model format can either be onnx or ort. One idea is to do the same for ort format optimizer model. This needs some investigation. - Changes to the offline tooling to generate ort format training artifacts. - End-to-end training example showcasing the use of the minimal training build. - Add support for export model for inferencing in a minimal build.
3.9 KiB
ORT Flatbuffer Schemas
This directory contains the ORT file format schema and the generated C++ header file for the ORT file format. This directory also contains the on-device training checkpoint schema and the corresponding auto generated C++ header file for the checkpoint format.
The ORT file format schema and the on-device training checkpoint schema uses the FlatBuffers serialization library.
Please do not directly modify the generated C++ header file for the ORT file format or for the training checkpoint file format, or the generated Python binding files.
The flatbuffers compiler (flatc) is built as part of an ONNX Runtime build. It is located in the _deps/flatbuffers-build/ subdirectory of the build output directory.
e.g.
- Windows Debug build
- \build\Windows\Debug_deps\flatbuffers-build\Debug\flatc.exe
- Linux Debug build
- /build/Linux/Debug/_deps/flatbuffers-build/flatc
It is possible to use another flatc as well, e.g., from a separate installation. Note that ONNX Runtime uses FlatBuffers 1.12.
To update the flatbuffers schemas and generated files:
-
Modify the ORT file format schema or training checkpoint schema.
-
Run compile_schema.py to generate the C++ and Python bindings.
python onnxruntime/core/flatbuffers/schema/compile_schema.py --flatc <path to flatc> -
Update the version history and record the changes. Changes made to the ORT file format schema warrants not only updating the ort format version, but also the checkpoint version since the checkpoint schema depends on the ort format schema.
ORT FB format version history
In ort_format_version.h, see IsOrtModelVersionSupported() for the supported versions and
kOrtModelVersion for the current version.
Version 1
History begins.
Initial support for FlatBuffers that includes Model support. Graph support including Attributes, Tensors, Tensor Sequences, Maps and Sequences. Constant initializers are also supported. Constant nodes are converted to constant initializers in the ORT format.
Version 2
Support for sparse initializers. Sparse intializers are stored within ORT FlatBuffers format, which includes sparse initializers converted from a Constant node attribute.
Version 3
Support for storing graph_doc_string field in Model (ORT FlatBuffers format).
Version 4
Update kernel def hashing to not depend on ordering of type constraint types (NOT BACKWARDS COMPATIBLE).
Version 5
Remove kernel def hashes and add KernelTypeStrResolver info to replace them (LIMITED BACKWARDS COMPATIBILITY). The change to the ORT format itself is not backwards compatibility-breaking, but ORT provides limited backwards compatibility for processing older models with missing KernelTypeStrResolver info. The motivation for this update is to support additional execution providers with statically registered kernels. The original approach of using kernel def hashes was not so extensible as it required the execution provider providing hashes to be enabled at model conversion time.
Version 6
Support for float 8 types. See Float stored in 8 bits for further details about their format and usage.
Checkpoint format version history
In checkpoint_version.h, see IsCheckpointVersionSupported() for the supported versions and
kCheckpointVersion for the current version.
Version 1
Initial support for the On-Device Training Checkpoint format. The format includes support for the ModuleState (stores the module parameters), OptimizerGroups (stores the optimizer states), and PropertyBag (stores custom user properties with support for int64, float and strings).