Fix nondeterministic Loop output binding with multiple state variables - #3041
Open
Amir Fathi (AmirF194) wants to merge 1 commit into
Open
Conversation
_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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
_translate_loop_stmtcomputesloop_state_varsonce (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 aslist(loop_state_vars | scan_outputs). A set union creates a new set, and its iteration order can differ fromloop_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_varsa list where it's first computed, and buildingoutputsby concatenation instead of a set union, so all four consumers walk the same order.With 5 loop-carried state variables, about a fifth of
PYTHONHASHSEEDvalues 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:
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 checkandruff format --checkclean on the changed files.Didn't reproduce the reporter's own
furthest_samplingmodel (needs torch); the test above hits the same code path with a smaller repro.