[ydshieh]

What does this PR do?

Fixes the flaky Wav2Vec2ConformerModelTest::test_onnx_export_dynamic (and the same risk for its siblings).

In short: Wav2Vec2ConformerForPreTraining has a discrete value selection (an argmax over 320 codes in its quantizer). When two candidates are almost tied, eager PyTorch and ONNX Runtime can pick different ones because of tiny float rounding differences. The flipped choice then changes the tensors computed from it, and the scalar codevector_perplexity no longer matches. The ONNX export itself is correct.

Fix (test-only): add Wav2Vec2ConformerForPreTraining and UniSpeechSatForPreTraining to EXPORT_SKIPS["onnx.exactness"], next to Wav2Vec2ForPreTraining and UniSpeechForPreTraining, which are already there for the same reason. These classes are still exported and run in ONNX Runtime; only the eager-vs-ONNX value comparison is dropped for them. The reason text is also corrected (see section 6).

Failures (Wav2Vec2ConformerModelTest::test_onnx_export_dynamic, CPU):

	main	this PR
40,000 random model / input draws (section 5)	15 failed (about 1 in 2,700)	–
those 15 failing draws, replayed through the real test (section 7)	15 / 15 failed	0 / 15 failed

---

Details (only for the record)

1. What the test checks

test_onnx_export_{dynamic,static} (from ExportTesterMixin in tests/exporters/test_export.py) does, for every class in all_model_classes:

1. build the tiny random model in eval mode with the tester's random inputs;
2. run eager PyTorch → reference outputs;
3. export with OnnxExporter and run the exported graph in ONNX Runtime on the same inputs;
4. check that the output names match, and that the values match eager with atol=rtol=1e-3, allowing up to 5% of the elements of a tensor to mismatch.

For Wav2Vec2ConformerForPreTraining (no labels / negative samples passed) there are 3 outputs:

output	shape in the test	goes through the argmax?
projected_states	(13, 14, 256)	no
projected_quantized_states	(13, 14, 256)	yes
codevector_perplexity	scalar	yes

2. The discrete value selection

In eval mode, Wav2Vec2ConformerGumbelVectorQuantizer.forward does (no Gumbel sampling in eval mode):

hidden_states = self.weight_proj(hidden_states)                   # logits: 320 codes per group
codevector_idx = hidden_states.argmax(dim=-1)                     # pick 1 code   <-- discrete selection
codevector_probs = hidden_states.new_zeros(...).scatter_(-1, codevector_idx.view(-1, 1), 1.0)  # one-hot
perplexity = self._compute_perplexity(codevector_probs, mask_time_indices)

With the test's tiny config this is 13 (batch) × 14 (frames) × 2 (groups) = 364 argmax calls per forward, each over 320 codes.

The flow after the argmax:

extract_features
  ├─ encoder → project_hid → projected_states                       (not affected)
  └─ quantizer: weight_proj → argmax → one-hot
        ├─ codevectors → project_q → projected_quantized_states      (affected)
        └─ code usage counts → exp(entropy) → codevector_perplexity  (affected)

codevector_perplexity is computed per group from how often each code is picked (the average of the one-hots over all frames): exp(entropy) of that distribution, summed over the 2 groups.

3. Why a code flips

Eager PyTorch and ONNX Runtime compute the same logits with different kernels, so the logits differ by a few float32 rounding steps (ulps). Measured on CPU, eager vs ONNX Runtime, max absolute difference per run (median over 300 runs):

tensor	max abs diff
conv feature extractor output	2.3e-6
extract_features (quantizer input)	1.9e-6
quantizer logits	9.5e-6

This is normal float noise, and it is the same in passing and failing runs (section 5). A code flips only when, in one of the 364 rows, the top two logits are closer than this noise. With random weights and random inputs this happens by chance; with logits around 10, 1 ulp is 9.5e-7.

4. What one flip does to the outputs

One flipped row means one frame gets a completely different codevector. Example from a captured failure, each output compared on its own:

output	result
projected_states	matches
projected_quantized_states	253 / 46592 elements (0.5%) mismatch, max diff 0.48: within the 5% allowance
codevector_perplexity	expected 92.1406, got 91.2174: fails (a scalar has no element-level allowance)

Over the 15 captured failures, one flip moved the perplexity by 0.09 to 0.92 (0.11% to 1.0% of its value). The allowed difference is 1e-3 + 1e-3 * |expected|, i.e. about 0.08 to 0.15.

5. Measurements

Done on CPU (CUDA_VISIBLE_DEVICES=), torch 2.15.0, A10G runner, with a script that reproduces the test's model / inputs / export for Wav2Vec2ConformerForPreTraining and saves the eager and ONNX Runtime intermediate tensors (logits, argmax, one-hot, perplexity, …) of every run.

• 40,000 random seeds were screened by the smallest top-1/top-2 logit gap in eager; the 273 with a gap < 1.5e-5 were exported and compared: 15 failures (about 1 in 2,700 runs overall on CPU).
• In all 15 failures, exactly 1 of the 364 rows flipped, and it was the row with the smallest gap. The eager gap between the two codes was 0 to 5.7e-6 (0 to 6 ulps). 3 of them were exact ties in eager.
• The noise, input scale, padding fraction and logit scale are the same for failing and passing runs. The only thing that distinguishes a failing run is a near-tie in the logits.
• 2 runs had a flip but still passed: the flip moved a row from a code used once to an unused code, which leaves the code usage counts, and so the perplexity, unchanged.

On CI the test fails more often than on CPU. The daily CI shows new failures with AssertionError: Scalars are not close! on single GPU (2026-10-23, 10-26, 11-03, 11-08) and multi GPU (2026-10-16, 10-22, 11-05, 11-08), and the test was P / F / P in 3 runs on a GPU runner. The GPU rate was not measured separately.

6. Why this fix

• onnx.exactness is the existing mechanism for "the export works but its outputs can't be compared to eager". Wav2Vec2ForPreTraining and UniSpeechForPreTraining already use it for this exact quantizer; Wav2Vec2ConformerForPreTraining and UniSpeechSatForPreTraining use the same quantizer and were missing.
  • UniSpeechSatForPreTraining has not failed on CI yet, but it is in UniSpeechSatModelTest.all_model_classes and has the same argmax, so it is added too.
  • The entries don't skip the test: these classes are still exported and run in ONNX Runtime, and their output names are still checked. The other classes of these models (Model, ForCTC, ForSequenceClassification, …) keep the full value comparison.
  • The scope applies to both test_onnx_export_dynamic and test_onnx_export_static; both have the same argmax.
• A higher tolerance doesn't fit: the mismatch is a different discrete choice, not a precision issue. Covering one flip would need rtol of about 1e-2 (observed) to 3e-2 (rough worst case), applied to every output of every class of the model.
• @is_flaky would keep the comparison, but needs a model-file override of the parameterized-generated test, and reruns a ~50 s test on each failure.
• The old reason was inaccurate. It said the outputs "come out of a Gumbel-softmax draw … and eager does not agree with itself either". In eval mode (as in the test) there is no Gumbel draw: eager is deterministic, and the mismatch comes from the argmax near-tie. The reason is updated, and the 3 other entries point to it.

The same argmax instability is already documented in these model test files: test_batching_equivalence is marked @is_flaky in wav2vec2, wav2vec2_conformer and unispeech, with the reason "The codevector_idx computed with argmax() … is not stable."

7. Verification

CPU (CUDA_VISIBLE_DEVICES=), RUN_SLOW=1, torch 2.15.0, onnxruntime 1.31.1, onnxscript 0.7.2.

Replay of captured failures through the real test. A failure happens about once in 2,700 CPU runs, so a before/after flakefinder run can't show the fix. Instead, a local pytest plugin (not part of this PR) seeds torch and the testers' global_rng with a seed captured from a failing run, right before the mixin builds Wav2Vec2ConformerForPreTraining. This recreates the exact failing model and inputs. All 15 failing seeds from section 5:

	main	this PR
15 failing seeds	15 failed, each on codevector_perplexity (e.g. seed 1371433937: expected 92.1406, got 91.2174)	15 passed

Normal runs with this PR (no seeding): test_onnx_export_{dynamic,static} for wav2vec2, wav2vec2_conformer, unispeech and unispeech_sat (all test classes): 12 passed.

Note

The test reports codevector_perplexity (not projected_quantized_states) because _check_outputs_close passes the whole output dict to one torch.testing.assert_close. That call checks the keys in sorted order and raises at the first mismatch, and the 5% allowance is applied only to that first error. If a tolerated mismatch came first, the outputs after it would not be compared at all. This doesn't change anything for this PR, but it could hide real mismatches in other models; it could be fixed separately by comparing output by output.

🤖 Generated with Claude Code

[ydshieh]

These last 2 are added in this PR.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[IlyasMoutawwakil]

yeh sure, i will later only skip the outputs where the near-tie problem arises and compare others like logits/projections.