Skip to content

Fix ASan OOM in QDQ Gemm transformer tests - #28797

Merged
tianleiwu merged 1 commit into
mainfrom
tlwu/20260604/fix_asan_oom
Jun 4, 2026
Merged

Fix ASan OOM in QDQ Gemm transformer tests#28797
tianleiwu merged 1 commit into
mainfrom
tlwu/20260604/fix_asan_oom

Conversation

@tianleiwu

Copy link
Copy Markdown
Contributor

PR: Fix ASan OOM in QDQ Gemm transformer tests

Description

PR #28131 ("Reject QDQ Gemm→QGemm fusion when alpha != 1 with bias") added alpha_not_one
coverage to the QDQ Gemm fusion tests. This multiplied the number of TransformerTester
session builds inside the already-large Gemm_U8U8U8 test matrix and pushed the
AddressSanitizer (ASan) build of onnxruntime_test_all over its allocator limit, causing the
windows_x64_asan CI to fail with AddressSanitizer: Out of memory. The process has exhausted 8192MB for size class 8192. This PR isolates the alpha != 1 coverage into small, dedicated
tests so the peak memory of any single test is reduced.

Summary of Changes

File Change
onnxruntime/test/optimizer/qdq_transformer_test.cc Added an opset_version parameter to QDQTransformerGemmTests (default 0 = run opsets 12/18/19); replaced the three hardcoded TransformerTester calls with a loop over the selected opset(s); removed the inline alpha_not_one block from the templated QDQTransformerGemmTests(); added a dedicated TEST(QDQTransformerTests, Gemm_AlphaNotOne_U8U8U8) that runs only uint8/uint8/uint8 at opset 19.
onnxruntime/test/optimizer/qdq_transformer_fastmath_test.cc Same refactor for the fastmath variant; added TEST(QDQTransformerTests, Gemm_AlphaNotOne_U8U8U8_FastMath) running only uint8/uint8/uint8 at opset 19.

The net effect is that the incremental alpha_not_one work added by #28131 drops from 24
session builds (4 alpha variants × 3 opsets, in each of the regular and fastmath files) to 8,
and is no longer part of the large Gemm_U8U8U8 matrix test — directly lowering the peak
memory consumed in a single test.

Testing

  • Build with --enable_address_sanitizer (the windows_x64_asan configuration) and run
    onnxruntime_test_all; confirm QDQTransformerTests.Gemm_AlphaNotOne_U8U8U8,
    QDQTransformerTests.Gemm_AlphaNotOne_U8U8U8_FastMath, and QDQTransformerTests.Gemm_U8U8U8
    pass and the suite no longer hits the ASan OOM.
  • Fusion behavior is unchanged: the same alpha != 1 rejection logic is still exercised, just
    with a narrower datatype/opset footprint.

Motivation and Context

The ASan failure is the sanitizer's internal allocator size-class limit (8 GB for
size class 8192), not a runner RAM cap that can simply be raised. Loosening it via
ASAN_OPTIONS quarantine tuning would weaken the sanitizer's bug-detection guarantees, so the
fix targets the test's memory footprint instead.

Options considered

  1. --test_parallel (reduce CTest concurrency). Lower the parallelism in the ASan workflow
    (e.g., --test_parallel 1) so fewer test binaries run concurrently. This only addresses
    cumulative/overlapping process memory; it does not reduce the peak memory of a single
    test, and it slows the CI down for every run. Rejected as a blunt, non-durable workaround.

  2. Shard the ASan tests. Split onnxruntime_test_all into N gtest shards
    (GTEST_TOTAL_SHARDS / GTEST_SHARD_INDEX) so the ASan allocator resets between shards.
    This helps with cumulative growth across the whole binary, but it still does not reduce
    the peak memory of any individual test — if one test alone approaches the limit, sharding
    the binary will not help. Rejected for the same root-cause reason.

  3. Break the test into smaller tests (chosen). Isolate the alpha != 1 coverage into
    dedicated tests that run a single datatype (uint8/uint8/uint8) at a single opset (19), and
    remove the alpha cases from the large Gemm_U8U8U8 matrix. This reduces the work done in
    the heaviest single test and addresses the peak-memory problem at its source while keeping
    the same fusion behavior under test.

Reference: PR #28131 (merge commit 585273033e).

Checklist

  • Tests added/updated
  • Documentation updated (if applicable)
  • No breaking changes (test-only change)
  • CI passes

@yuslepukhin

Copy link
Copy Markdown
Contributor

LGTM. The fix is minimal, well-targeted at the actual root cause (per-test peak), and preserves the substantive coverage of the alpha != 1 rejection path. The only material trade-off is opset breadth for the alpha-not-one case, which is acceptable. I'd suggest the optional in-source comment about the rationale before merging.

Suggestions (optional, non-blocking)

Add a one-line comment above each new TEST explaining why coverage is restricted to opset 19, e.g. // Isolated from Gemm_U8U8U8 to keep ASan peak memory under the size-class limit; alpha != 1 rejection is opset-agnostic. That preserves the rationale in-tree for future maintainers (the PR description will not be visible from the source).

If the team wants to retain alpha_not_one coverage at opsets 12 and 18, a follow-up could add Gemm_AlphaNotOne_U8U8U8_Opset12 / _Opset18 tests; each would still be small. Not required for this PR.

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 refactors the QDQ Gemm transformer test matrix to reduce peak per-test memory (specifically addressing ASan OOM on windows_x64_asan) by moving alpha != 1 coverage out of the already-large Gemm_U8U8U8 matrix and into dedicated, narrower tests.

Changes:

  • Added an opset_version parameter to the QDQTransformerGemmTests helpers and replaced hard-coded TransformerTester calls with a loop over the selected opset(s).
  • Removed inline alpha_not_one coverage from the templated Gemm matrix tests and introduced dedicated Gemm_AlphaNotOne_U8U8U8 tests (regular + fastmath) at opset 19 only.
  • Preserved existing fusion expectations while reducing the number of session builds that happen inside the heaviest test.

Reviewed changes

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

File Description
onnxruntime/test/optimizer/qdq_transformer_test.cc Refactors Gemm test helper to loop opsets and adds a dedicated alpha-not-one Gemm test to reduce peak memory.
onnxruntime/test/optimizer/qdq_transformer_fastmath_test.cc Mirrors the same opset-loop refactor for fastmath tests and adds a dedicated alpha-not-one fastmath Gemm test.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread onnxruntime/test/optimizer/qdq_transformer_test.cc
Comment thread onnxruntime/test/optimizer/qdq_transformer_fastmath_test.cc
@tianleiwu
tianleiwu enabled auto-merge (squash) June 4, 2026 23:01
@tianleiwu
tianleiwu merged commit d43a48b into main Jun 4, 2026
86 of 87 checks passed
@tianleiwu
tianleiwu deleted the tlwu/20260604/fix_asan_oom branch June 4, 2026 23:43
GopalakrishnanN added a commit that referenced this pull request Jun 5, 2026
…in session_state_test (#28799)

## Description

Fixes a latent bug and reduces unnecessary resource usage in
`session_state_test.cc`.

### Thread-pool parameter initialization

`SessionStateTestP.TestInitializerProcessing` contained a
self-assignment that silently dropped the test's intended thread count:

```cpp
OrtThreadPoolParams to;
to.thread_pool_size = to.thread_pool_size;  // no-op: leaves thread_pool_size = 0
```

`thread_pool_size = 0` means "use all physical cores (or half the
logical cores)". Across the 8 parametrized instances this spins up large
intra-op thread pools regardless of the `thread_count` parameter, adding
avoidable thread/memory overhead — which is especially costly under the
AddressSanitizer build.

This PR:
- Corrects the assignment to use the intended parameter:
`to.thread_pool_size = param.thread_count;`
- Value-initializes all local `OrtThreadPoolParams` instances
(`OrtThreadPoolParams to{};`) for consistency and to avoid relying on
implicit defaults.

### Deduplicate schema registration

`PrePackingTest` was registered into the global ONNX schema registry in
two places (`SessionStatePrepackingTest` and the
`SessionStateTestSharedInitalizersWithPrePacking` fixture), producing
duplicate-registration warnings when both run in the same process. The
registration is now centralized behind a `std::call_once` helper so the
schema is registered exactly once.

## Motivation and Context

While investigating an ASan `onnxruntime_test_all` OOM,
`SessionStatePrepackingTest` showed up as the abort site (the
process-wide ASan allocator aborts in whichever test is allocating when
the limit is hit). The dominant memory driver was addressed separately
in #28797 (QDQ Gemm transformer tests). This PR cleans up the genuine,
independent issues found in `session_state_test.cc` during that
investigation: the self-assignment bug and the duplicate schema
registration.

## Testing

All affected tests pass locally (RelWithDebInfo):

```
[==========] 19 tests from 4 test suites ran.
[  PASSED  ] 19 tests.
```

Covers `SessionStateTestP`, `SessionStatePrepackingTest`,
`SessionStateAddGetKernelTest`, and
`SessionStateTestSharedInitalizersWithPrePacking`.

---------

Co-authored-by: Gopalakrishnan Nallasamy <gnallasamy@microsoft.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants