onnxruntime/onnxruntime/core/flatbuffers/schema/ort.fbs
Scott McKay 9372e9a0a3
Support >2GB of Tensor data in training checkpoint (#20077)
### Description
<!-- Describe your changes. -->
Add ability to store initializer data in an external file.
Update training checkpoint code to use external file if data > ~2GB.

I don't see a way for the flatbuffers 64-bit offsets to be used, as they
don't support storing 'table' types with 64-bit offsets (and our Tensor
is a 'table' type not a simple struct).


0cfb7eb80b/tests/64bit/test_64bit.fbs (L38-L39)

Allowing a Tensor to have its raw_data in an external file should
hopefully work with the least friction. As it's an extra field it's
backwards compatible.

Please feel free to suggest alternative approaches. 

Side note: the diffs in the generated *.fbs.h files are unexpectedly
large. Maybe they weren't re-generated when the new flatbuffers version
was checked in. I updated by running:
`python .\compile_schema.py -f <build output
dir>\_deps\flatbuffers-build\Debug\flatc.exe`
from onnxruntime\core\flatbuffers\schema which I thought was the correct
way but maybe that's out of date.

I think you can ignore all the diffs in the generated files and just
worry about the changes to the .fbs files in
onnxruntime/core/flatbuffers/schema. Basically start at the bottom of
the files changed and work up as all the 'real' diffs are there.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

---------

Co-authored-by: carzh <wolfivyaura@gmail.com>
2024-04-22 15:17:43 -07:00

325 lines
6.1 KiB
Text

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
namespace onnxruntime.fbs;
// Attribute
enum AttributeType : int32 {
UNDEFINED = 0,
FLOAT = 1,
INT = 2,
STRING = 3,
TENSOR = 4,
GRAPH = 5,
FLOATS = 6,
INTS = 7,
STRINGS = 8,
TENSORS = 9,
GRAPHS = 10,
SPARSE_TENSOR = 11,
SPARSE_TENSORS = 12,
}
// Shape
table Shape {
dim:[Dimension];
}
table Dimension {
value:DimensionValue;
denotation:string;
}
enum DimensionValueType : int8 {
UNKNOWN = 0,
VALUE = 1,
PARAM = 2,
}
table DimensionValue {
dim_type:DimensionValueType;
dim_value:int64;
dim_param:string;
}
// Tensor
enum TensorDataType : int32 {
UNDEFINED = 0,
FLOAT = 1,
UINT8 = 2,
INT8 = 3,
UINT16 = 4,
INT16 = 5,
INT32 = 6,
INT64 = 7,
STRING = 8,
BOOL = 9,
FLOAT16 = 10,
DOUBLE = 11,
UINT32 = 12,
UINT64 = 13,
COMPLEX64 = 14,
COMPLEX128 = 15,
BFLOAT16 = 16,
// Float 8 types. See https://onnx.ai/onnx/technical/float8.html.
FLOAT8E4M3FN = 17,
FLOAT8E4M3FNUZ = 18,
FLOAT8E5M2 = 19,
FLOAT8E5M2FNUZ = 20,
}
table TensorTypeAndShape {
elem_type:TensorDataType;
shape:Shape;
}
table MapType {
key_type:TensorDataType;
value_type:TypeInfo;
}
table SequenceType {
elem_type:TypeInfo;
}
// Node
enum NodeType : int32 {
Primitive = 0,
Fused = 1,
}
struct EdgeEnd {
node_index:uint32;
src_arg_index:int32;
dst_arg_index:int32;
}
table NodeEdge {
node_index:uint32;
input_edges:[EdgeEnd];
output_edges:[EdgeEnd];
}
table Node {
name:string;
doc_string:string;
domain:string;
since_version:int32;
index:uint32;
op_type:string;
type:NodeType;
execution_provider_type:string;
inputs:[string];
outputs:[string];
attributes:[Attribute];
input_arg_counts:[int32];
implicit_inputs:[string];
}
// ValueInfo
table ValueInfo {
name:string;
doc_string:string;
type:TypeInfo;
}
// TODO add support of SparseTensor, Opaque if needed
union TypeInfoValue {
tensor_type:TensorTypeAndShape,
sequence_type:SequenceType,
map_type:MapType,
}
table TypeInfo {
denotation:string;
value:TypeInfoValue;
}
// OpSetId
table OperatorSetId {
domain:string;
version:int64;
}
// For simplicity, we will have only two data fields
// - string_data for string
// - raw_data for all other types
table Tensor {
name:string;
doc_string:string;
dims:[int64];
data_type:TensorDataType;
raw_data:[uint8];
// string_data is least used
string_data:[string];
// offset into external data file to allow data >2GB to be handled. not used for string data.
// an external file writer/reader needs to be provided when serializing.
// int64 (vs uint64) so we can explicitly set to -1 when not used.
external_data_offset:int64 = -1;
}
table SparseTensor {
values:Tensor;
indices:Tensor;
dims:[int64];
}
table Attribute {
name:string;
doc_string:string;
type:AttributeType;
f:float32;
i:int64;
s:string;
t:Tensor;
g:Graph;
floats:[float32];
ints:[int64];
strings:[string];
tensors:[Tensor];
graphs:[Graph];
}
// runtime optimizations
/// nodes to consider for a runtime optimization
/// see corresponding type in onnxruntime/core/graph/runtime_optimization_record.h
table NodesToOptimizeIndices {
node_indices:[uint32];
num_inputs:uint32;
num_outputs:uint32;
has_variadic_input:bool;
has_variadic_output:bool;
num_variadic_inputs:uint32;
num_variadic_outputs:uint32;
}
/// deprecated: no longer using kernel def hashes
table DeprecatedNodeIndexAndKernelDefHash {
node_index:uint32;
kernel_def_hash:uint64;
}
/// a single runtime optimization
/// see corresponding type in onnxruntime/core/graph/runtime_optimization_record.h
table RuntimeOptimizationRecord {
action_id:string;
nodes_to_optimize_indices:NodesToOptimizeIndices;
produced_nodes:[DeprecatedNodeIndexAndKernelDefHash] (deprecated);
produced_op_ids:[string];
}
table RuntimeOptimizationRecordContainerEntry {
optimizer_name:string (key);
runtime_optimization_records:[RuntimeOptimizationRecord];
}
table RuntimeOptimizations {
/// mapping from optimizer name to [RuntimeOptimizationRecord]
records:[RuntimeOptimizationRecordContainerEntry];
}
table Graph {
initializers:[Tensor];
node_args:[ValueInfo];
nodes:[Node];
max_node_index:uint32;
node_edges:[NodeEdge];
inputs:[string];
outputs:[string];
sparse_initializers:[SparseTensor];
runtime_optimizations:RuntimeOptimizations;
}
table StringStringEntry {
key:string;
value:string;
}
table Model {
ir_version:int64;
opset_import:[OperatorSetId];
producer_name:string;
producer_version:string;
domain:string;
model_version:int64;
doc_string:string;
graph:Graph;
graph_doc_string:string;
metadata_props:[StringStringEntry];
}
/// deprecated: no longer using kernel def hashes
table DeprecatedKernelCreateInfos {
node_indices:[uint32];
kernel_def_hashes:[uint64];
}
/// deprecated: no longer using kernel def hashes
table DeprecatedSubGraphSessionState {
// graph_id can be used to binary search DeprecatedSubGraphSessionState in
// DeprecatedSessionState.sub_graph_session_states
graph_id:string (key);
session_state:DeprecatedSessionState;
}
/// deprecated: no longer using kernel def hashes
table DeprecatedSessionState {
kernels:DeprecatedKernelCreateInfos;
sub_graph_session_states:[DeprecatedSubGraphSessionState];
}
enum ArgType : int8 {
INPUT = 0,
OUTPUT = 1,
}
table ArgTypeAndIndex {
arg_type:ArgType;
index:uint32;
}
table KernelTypeStrArgsEntry {
kernel_type_str:string (key);
args:[ArgTypeAndIndex];
}
table OpIdKernelTypeStrArgsEntry {
op_id:string (key);
kernel_type_str_args:[KernelTypeStrArgsEntry];
}
table KernelTypeStrResolver {
op_kernel_type_str_args:[OpIdKernelTypeStrArgsEntry];
}
table InferenceSession {
// This is the ORT format model version
// The version number is defined as kOrtModelVersion in <repo root>/onnxruntime/core/flatbuffers/ort_format_version.h
ort_version:string;
model:Model;
session_state:DeprecatedSessionState (deprecated);
kernel_type_str_resolver:KernelTypeStrResolver;
}
root_type InferenceSession;
file_identifier "ORTM";