Skip to content

Refactor warmup logic and generator mock in OnnxDiscrepancyCheck#2568

Merged
tadani3 merged 3 commits into
mainfrom
tommasoadani/genai_discrepancy_bugfix
Jul 16, 2026
Merged

Refactor warmup logic and generator mock in OnnxDiscrepancyCheck#2568
tadani3 merged 3 commits into
mainfrom
tommasoadani/genai_discrepancy_bugfix

Conversation

@tadani3

@tadani3 tadani3 commented Jul 13, 2026

Copy link
Copy Markdown
Contributor

Summary

Fix a warmup pollution bug in OnnxDiscrepancyCheck.compare_generation
(olive/passes/onnx/discrepancy_check.py) that caused a systematic off-by-one
misalignment in the ONNX Runtime GenAI generated-token list.

Bug

The GenAI warmup loop (introduced in fe69db6) reused the genai_tokens
accumulator across both the warmup and the real generation run. The generator
is correctly re-created for the real run, but the accumulator was not reset, so
warmup-generated tokens were already present in genai_tokens when the real run
started.

Because generation is deterministic, this shifted every generated token by one
position and produced misleading metrics:

  • first_token_matchesTrue (warmup token equals the real first token)
  • second_token_matchesFalse (real first token lands in the second slot)
  • longest_common_token_sequence → collapsed to 1

The warmup break condition (len(genai_tokens) >= 3) also counted prompt tokens,
under-warming on prompts of 3 or more tokens.

Fix

  • Use a throwaway warmup_tokens accumulator during warmup so genai_tokens is
    never touched until the real run — structurally mirroring how the transformers
    side already discards its warmup output.
  • Change the warmup break condition to len(warmup_tokens) >= genai_prompt_token_count + 1
    so it counts only generated tokens.

The real generation block immediately following was already correct and is
unchanged.

Tests

  • Updated test/passes/onnx/test_discrepancy_check.py TestCompareGeneration
    mocks so each og.Generator call (warmup + real run) replays the token stream
    from the start via an og.Generator.side_effect factory. This models real
    deterministic re-generation, where warmup and the real run produce the same
    tokens and the warmup output is discarded. Previously the mocks shared a single
    call_count across both generators, which encoded the buggy behavior.
  • The two append_tokens assertions now target a shared mock so the "prompt is
    seeded for both warmup and real run" checks still hold.
  • test_compare_generation_with_zero_max_new_tokens was left unchanged — its
    generator reports is_done immediately and never appended warmup tokens, so it
    was unaffected by the bug.

Checklist before requesting a review

  • Add unit tests for this change.
  • Make sure all tests can pass.
  • Update documents if necessary.
  • Lint and apply fixes to your code by running lintrunner -a
  • Is this a user-facing change? If yes, give a description of this change to be included in the release notes.

Release notes

Fixed a bug in the OnnxDiscrepancyCheck pass where GenAI warmup tokens polluted
the generated-token comparison, causing incorrect first_token_matches,
second_token_matches, and longest_common_token_sequence results.

(Optional) Issue link

@tadani3
tadani3 requested review from Copilot and xadupre July 13, 2026 17:23
@tadani3 tadani3 self-assigned this Jul 13, 2026
@tadani3 tadani3 added the bug Something isn't working label Jul 13, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR fixes a GenAI warmup “token pollution” bug in OnnxDiscrepancyCheck.compare_generation that could misalign generated-token comparisons (systematically shifting GenAI tokens by one) and thereby corrupt first_token_matches, second_token_matches, and longest_common_token_sequence results.

Changes:

  • Refactored GenAI warmup to append into a throwaway warmup_tokens list instead of genai_tokens, preventing warmup output from affecting the real run.
  • Adjusted the warmup break condition to be based on generated-token count (via genai_prompt_token_count + 1 threshold) rather than a fixed total-token length.
  • Updated unit tests to mock og.Generator with a side_effect factory so each generator instance (warmup + real) replays the deterministic token stream from the start, matching real behavior.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File Description
olive/passes/onnx/discrepancy_check.py Prevents warmup-generated tokens from being appended into the “real run” accumulator and fixes warmup stopping criteria.
test/passes/onnx/test_discrepancy_check.py Updates generator mocking to model deterministic re-generation per generator instance and keeps prompt-seeding assertions valid.

@tadani3
tadani3 merged commit 7493634 into main Jul 16, 2026
12 checks passed
@tadani3
tadani3 deleted the tommasoadani/genai_discrepancy_bugfix branch July 16, 2026 18:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants