microsoft / microsoft/onnxruntime
CompileModel drops outer-graph initializers consumed only from a subgraph, then rejects its own output
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Describe the issue
`CompileModel` / `OrtCompileApi` fails on any model where a subgraph (`If` / `Loop` / `Scan` body)
reads an initializer of an enclosing graph that no node of that enclosing graph reads itself.
The ONNX IR spec allows this explicitly:
> A node input in a nested subgraph MAY refer to names introduced in outer graphs (as node
> outputs, graph inputs, or **graph initializers**).
> -- https://github.com/onnx/onnx/blob/main/docs/IR.md
`onnx.checker.check_model(..., full_check=True)` accepts such a model, `InferenceSession` builds
and runs it correctly, and `SessionOptions.optimized_model_filepath` round-trips it intact. Only
the compile API rejects it, and the error describes the model ORT itself generated, not the input.
**Root cause.** `CreateEpContextModel` in `onnxruntime/core/framework/graph_partitioner.cc`
([v1.24.4](https://github.com/microsoft/onnxruntime/blob/v1.24.4/onnxruntime/core/framework/graph_partitioner.cc#L921-L928),
unchanged on `main`) copies an initializer only when a `NodeArg` of that name already exists in
the new graph:
```cpp
// handle initializers
for (const auto& [name, _] : graph.GetAllInitializedTensors()) {
if (ep_graph.GetNodeArg(name) != nullptr) {
graph_utils::MakeInitializerCopyIfNotExist(graph, ep_graph, name);
}
}
ORT_RETURN_IF_ERROR(ep_graph.Resolve());
```
`ep_graph.AddNode(node)` creates `NodeArg`s for the node's own inputs and outputs. A name
referenced only from inside the node's subgraph attribute gets none, so the initializer is skipped
and the following `Resolve()` fails on the dangling outer-scope reference.
The information needed is already on the source node: `Node::ImplicitInputDefs()` lists exactly
the `NodeArg`s a node's subgraphs consume from outer scope, and it is what `Graph::IsOuterScopeValue`
uses. Copying initializers named by the implicit input defs as well would fix it.
This looks like the same class of bug as #23043, which was `BeamSearch` subgraph setup missing the
same information, and which was fixed.
### To reproduce
The script builds two models that differ only in whether one extra outer-graph node also reads the
initializer, and runs the checker, a session, and the compiler against each.
```python
import numpy as np
import onnx
import onnxruntime as ort
from onnx import TensorProto, helper
def build(shared_axes_only_in_subgraph: bool) -> onnx.ModelProto:
# then-branch has no inputs of its own: it reads "x" and "axes" from the
# enclosing graph, which ONNX IR explicitly allows for initializers too.
then_g = helper.make_graph(
[helper.make_node("Unsqueeze", ["x", "axes"], ["t"], name="Unsqueeze_0")],
"then_branch", [],
[helper.make_tensor_value_info("t", TensorProto.FLOAT, [1, 1, 4])])
else_g = helper.make_graph(
[helper.make_node("Constant", [], ["e"], name="Sentinel",
value=onnx.numpy_helper.from_array(
np.zeros((1, 1, 4), np.float32), "zero"))],
"else_branch", [],
[helper.make_tensor_value_info("e", TensorProto.FLOAT, [1, 1, 4])])
nodes = [
helper.make_node("Shape", ["x"], ["s"], name="Shape"),
helper.make_node("Gather", ["s", "idx"], ["d"], name="Gather", axis=0),
helper.make_node("Equal", ["d", "four"], ["cond"], name="Equal"),
]
if not shared_axes_only_in_subgraph:
# one extra outer-graph reference to "axes" is enough to make it survive
nodes.append(helper.make_node("Identity", ["axes"], ["unused"], name="Echo"))
nodes.append(helper.make_node("If", ["cond"], ["y"], name="If0",
then_branch=then_g, else_branch=else_g))
g = helper.make_graph(
nodes, "outer",
[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, "w"])],
[helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 4])],
initializer=[
onnx.numpy_helper.from_array(np.array(1, np.int64), "idx"),
onnx.numpy_helper.from_array(np.array(4, np.int64), "four"),
onnx.numpy_helper.from_array(np.array([0], np.int64), "axes"),
])
m = helper.make_model(g, opset_imports=[helper.make_opsetid("", 21)], ir_version=10)
onnx.checker.check_model(m, full_check=True)
return m
print("onnxruntime", ort.__version__, "/ onnx", onnx.__version__)
for only_in_subgraph in (True, False):
tag = "read only from the subgraph" if only_in_subgraph else "also read by an outer node"
path = f"repro_{only_in_subgraph}.onnx"
onnx.save(build(only_in_subgraph), path)
sess = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
out = sess.run(None, {"x": np.arange(4, dtype=np.float32).reshape(1, 4)})[0]
print(f"\n'axes' {tag}")
print(" onnx.checker(full_check=True): pass")
print(f" InferenceSession: pass, output {out.ravel()}")
opts = ort.SessionOptions()
opts.add_provider("CPUExecutionProvider", {})
try:
ort.ModelCompiler(opts, path, embed_compiled_data_into_model=True) \
.compile_to_file(f"repro_{only_in_subgraph}_ctx.onnx")
print(" CompileModel: pass")
except Exception as exc:
print(f" CompileModel: FAIL\n {str(exc)}")
```
Output:
```
onnxruntime 1.28.0 / onnx 1.21.0
'axes' read only from the subgraph
onnx.checker(full_check=True): pass
InferenceSession: pass, output [0. 1. 2. 3.]
CompileModel: FAIL
[ONNXRuntimeError] : 10 : INVALID_GRAPH : This is an invalid model. In Node, ("If0", If, "", -1) : ("cond": tensor(bool),) -> ("y": tensor(float),) , Error Nodes in a graph must be topologically sorted, however input 'axes' of node:
name: Unsqueeze_0 OpType: Unsqueeze
is not output of any previous nodes.
'axes' also read by an outer node
onnx.checker(full_check=True): pass
InferenceSession: pass, output [0. 1. 2. 3.]
CompileModel: pass
```
Neither the EP nor the optimization level matters: the failure is identical on the CPU EP and on
DirectML, and at every graph optimization level from `ORT_DISABLE_ALL` (the compile API's default)
through `ORT_ENABLE_ALL`.
### Urgency
Not blocking, but it rules the compile API out for a whole class of models. Outer-scope
initializers are how sibling subgraphs share weights. The model I hit this with keeps five
statically shaped resolution branches behind chained `If` nodes and stores ~99% of its weight bytes
once in the outer graph, which makes the file a fifth of the size it would otherwise have. It
cannot be compiled at all today.
### Platform
Windows
### OS Version
Windows 11 Pro, build 10.0.26200
### ONNX Runtime Installation
Released Package
### ONNX Runtime Version or Commit ID
1.28.0, also reproduced on 1.24.4
### ONNX Runtime API
Python
### Architecture
X64
### Execution Provider
Default CPU, DirectML
### Execution Provider Library Version
DirectML 1.15.4
Contributor guide
Research direction
Inspect CreateEpContextModel in onnxruntime/core/framework/graph_partitioner.cc, especially its initializer handling, and compare it with Node::ImplicitInputDefs(). Run the supplied Python reproducer against models where the initializer is read only by a subgraph and where an outer node also reads it; done means both models compile successfully without changing inference behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 75/100