Skip to content

Fix nondeterministic Loop output binding with multiple state variables - #3041

Open
Amir Fathi (AmirF194) wants to merge 1 commit into
microsoft:mainfrom
AmirF194:fix/2203-nondeterministic-loop-output-order
Open

Fix nondeterministic Loop output binding with multiple state variables#3041
Amir Fathi (AmirF194) wants to merge 1 commit into
microsoft:mainfrom
AmirF194:fix/2203-nondeterministic-loop-output-order

Conversation

@AmirF194

Copy link
Copy Markdown

_translate_loop_stmt computes loop_state_vars once (vars_def_in_loop.intersection(exposed_uses | live_out)) and reuses that same set object for the loop body's parameters, the loop body's outputs, and the Loop node's loop-carried inputs, so those three always agree with each other.

outputs, the list used to bind each Python variable name back to a result value, is built separately as list(loop_state_vars | scan_outputs). A set union creates a new set, and its iteration order can differ from loop_state_vars's own order under CPython's hash randomization even though it holds the same elements. ONNX's Loop op matches inputs to outputs positionally, so when the two orders diverge, the converter binds a Python variable to the wrong computed value.

Fixed by making loop_state_vars a list where it's first computed, and building outputs by concatenation instead of a set union, so all four consumers walk the same order.

With 5 loop-carried state variables, about a fifth of PYTHONHASHSEED values reproduced the swap for me, which matches the "sometimes correct, sometimes wrong" symptom in the report. furthest_sampling's own while loop has 4.

Fixes #2203

Verification:

  • Added tests/loop_test.py::test_loop_state_var_output_order_matches_eager, a 5-state-variable loop run under 3 fixed hash seeds (3, 4, 8) that compares eager output to the exported model's onnxruntime output. On main it's red on all three seeds; on this branch it's green.
  • pytest tests/loop_test.py onnxscript/_internal/converter_test.py: 56 passed, same pre-existing skip/xfail/xpass on both sides.
  • ruff check and ruff format --check clean on the changed files.

Didn't reproduce the reporter's own furthest_sampling model (needs torch); the test above hits the same code path with a smaller repro.

_translate_loop_stmt binds each returned Python variable to a Loop node
output by walking `outputs = list(loop_state_vars | scan_outputs)`, a set
built separately from `loop_state_vars` itself. Under CPython's hash
randomization the two can iterate in different orders even with identical
elements, and ONNX matches Loop inputs to outputs positionally, so the
converter could bind a variable to the wrong computed value.

Make loop_state_vars a list where it's first computed and build outputs
by concatenation, so the loop body's parameters, its own outputs, the
Loop node's inputs, and the output binding all walk the same order.

Fixes microsoft#2203
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Development

Successfully merging this pull request may close these issues.

The calculation results of modelProto exported by onnxscript are variable

1 participant