ort-transformers-gpu-pytest

name

ort-transformers-gpu-pytest

description

Run the ONNX Runtime transformers Python tests (onnxruntime/test/python/transformers) against a GPU wheel, and confirm real cuDNN/flash SDPA dispatch. Use when a transformers pytest fails with ModuleNotFoundError onnxruntime.capi, when torch's bundled CUDA/cuDNN libs shadow the ORT-built ones (wrong-version dispatch or load errors), or when you must PROVE a test exercised the cuDNN SDPA decode tier instead of silently skipping/falling back to MATH.

Running ONNX Runtime transformers GPU pytest (dispatch-verified)

Reusable, hard-won knowledge for running the Python transformers tests under onnxruntime/test/python/transformers/ against a GPU-built wheel, on a box where torch is also installed. Three gotchas below will each cost you an hour if rediscovered. See the ort-test skill for the general test taxonomy and the false-green modes; this skill is the GPU + transformers + SDPA-dispatch specialization.

1. NEUTRAL-CWD gotcha — never run pytest from the repo root

Symptom:

ModuleNotFoundError: No module named 'onnxruntime.capi'
# or an AttributeError deep inside onnxruntime import

Cause: the repo root contains a source directory ./onnxruntime/ (the C++/Python source tree). When pytest runs with the repo root on sys.path[0], import onnxruntime resolves to that source package — which has no compiled capi extension — instead of the installed wheel in your venv's site-packages. The source dir shadows the wheel.

Fix: run pytest from a neutral, private working directory — a fresh mktemp -d, not the repo root and not a bare shared /tmp — so the shadowing source dir is not on the path, and point at the test file by absolute path. Put the transformers test-helper dir on PYTHONPATH so shared helpers still import.

WORKDIR=$(mktemp -d); cd "$WORKDIR"       # NEUTRAL + PRIVATE cwd — NOT repo root, NOT /tmp
export PYTHONPATH=/abs/repo/onnxruntime/test/python/transformers
python -m pytest /abs/repo/onnxruntime/test/python/transformers/<file>.py -v

Why mktemp -d and not a bare cd /tmp: pytest prepends the cwd to sys.path, and Python auto-imports sitecustomize.py/usercustomize.py from it at startup — so on a shared box a co-tenant's planted /tmp/sitecustomize.py would execute as arbitrary code in your test process. A fresh per-run mktemp -d (private, 0700) keeps the neutral-cwd source-shadow protection while removing that injection vector.

Do not cd into the repo and run pytest onnxruntime/test/... — that reintroduces the shadowing. (This is the Python analogue of the C++ "run from the build output dir" rule in ort-test.)

2. LD_PRELOAD lib-pinning gotcha — torch's CUDA/cuDNN shadow ORT's

Symptom: any of —

• the SDPA decode tier silently routes to MATH instead of CUDNN_FLASH_ATTENTION (wrong-version cuDNN loaded), or
• libcudnn.so.9: cannot open shared object file / undefined symbol / cuDNN version mismatch errors at first CUDA op, or
• ORT loads a different CUDA runtime than it was built against.

Cause: a pip-installed torch ships its own bundled CUDA runtime + cuDNN (e.g. cu124 → CUDA 12.4 / cuDNN 9.1) under site-packages/nvidia/*/lib. If ORT was built against a different CUDA/cuDNN (e.g. CUDA 12.9 / cuDNN 9.8), whichever set the dynamic loader resolves first wins. With torch imported (or its libs on the path), torch's older libs can shadow the ones ORT dlopens → wrong-version dispatch or load failure.

Fix: LD_PRELOAD the system CUDA runtime + cuDNN that ORT was built against so they are loaded first, and add their dirs to LD_LIBRARY_PATH. Activate the venv that has the ORT wheel. Concrete form used successfully (CUDA 12.9 + cuDNN 9.8; substitute your absolute lib paths):

WORKDIR=$(mktemp -d); cd "$WORKDIR"       # neutral + private (see §1)
source /abs/repo/.venv/bin/activate
export LD_PRELOAD=/abs/cuda12.9/lib64/libcudart.so.12:/abs/cudnn9.8/lib/libcudnn.so.9
export LD_LIBRARY_PATH=/abs/cuda12.9/lib64:/abs/cudnn9.8/lib
export PYTHONPATH=/abs/repo/onnxruntime/test/python/transformers
python -m pytest /abs/repo/onnxruntime/test/python/transformers/<file>.py -v

Notes:

• libcudart.so.12 is correct for both CUDA 12.4 and 12.9 (SONAME is major-only) — pinning the 12.9 file forces the right minor.
• torch.cuda still works fine under this preload — the bf16 IO-binding path that uses torch tensors + .data_ptr() runs correctly.
• Keep the two exports and the preload together; dropping LD_LIBRARY_PATH can still let a transitive dependency resolve against torch's copy.

3. Confirm REAL cuDNN SDPA dispatch (don't trust value-equality)

A numerically-correct result does not prove the cuDNN SDPA path ran — the kernel has a MATH fallback that produces the same answer (false-green mode 4 in ort-test). To prove the tier dispatched, observe ORT's routing rather than probing a version.

Observe-dispatch (the correct probe)

ORT's ONNX-domain Attention kernel emits a debug line when ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO=1 is set before the InferenceSession is created (the option is read once at session creation). AttentionKernelDebugInfo::Print emits a token of the form:

SdpaKernel=CUDNN_FLASH_ATTENTION      # or =MATH, =FLASH_ATTENTION, =EFFICIENT_ATTENTION (non-exhaustive)

Capture stdout across a single run() and parse it. Capture at the file-descriptor level, not contextlib.redirect_stdout: the SdpaKernel= line is written to native fd-1 from C++, which Python-level stdout redirection never intercepts — you would get dispatched=None, a silent false-negative. Mirror ORT's own _CaptureStdout (os.dup2 fd-1 to a temp file, run, restore, read it back):

os.environ["ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO"] = "1"   # BEFORE InferenceSession()

# FD-level capture (see onnxruntime's _CaptureStdout for the exact idiom):
saved_fd = os.dup(1)
tmp = tempfile.TemporaryFile()
os.dup2(tmp.fileno(), 1)            # redirect native fd-1
try:
    # ... create session, run once ...
finally:
    os.dup2(saved_fd, 1)           # restore fd-1
    os.close(saved_fd)
tmp.seek(0)
captured_text = tmp.read().decode()

m = re.search(r"SdpaKernel=(?P<kernel>[A-Z_]+)", captured_text)
dispatched = m.group("kernel") if m else None
assert dispatched == "CUDNN_FLASH_ATTENTION"

Caveat: re.search returns only the first SdpaKernel= token — correct for the single-node decode probe here. For a graph with multiple attention nodes use re.findall and check every token, or a later node's MATH fallback is masked by an earlier cuDNN hit.

Prefer this over reading torch.backends.cudnn.version() or any library-version check: a version probe reads torch's cuDNN, not the cuDNN ORT actually loaded/dispatched — that mismatch is a real trustworthiness bug. Observe-dispatch reads ORT's own routing decision, so it is correct across cuDNN versions with no hard-coded version table.

Non-skippable canary — ORT_TEST_REQUIRE_CUDNN_SDPA

Gating decode tests by observed dispatch has a failure mode: if the tier silently regresses (stops selecting cuDNN), the observation returns "not dispatched" and every decode test skips green, hiding the regression as all-green skips.

Close the hole with an env-gated canary: ORT_TEST_REQUIRE_CUDNN_SDPA=1. When set, the dispatch assertion becomes non-skippable — a MATH fallback / non-dispatch on the minimal known-good config FAILS LOUD instead of skipping. When unset (dev boxes, unsupported cuDNN) it falls back to the normal skip guard so it never false-alarms.

The variable is intended for an operator to export on a known-good GPU CI leg once one exists. Note that today no ONNX Runtime pipeline definition exports it (there is no Hopper+ GPU CI leg), so it has no effect in this project's CI and only matters for manual/local runs where a developer sets it explicitly — don't describe it in test docstrings as an enforcement that CI already applies.

def require_cudnn_sdpa():
    return os.environ.get("ORT_TEST_REQUIRE_CUDNN_SDPA") == "1"

# in the test:
enforce = require_cudnn_sdpa()
if not enforce and not cudnn_decode_supported(head_size):  # illustrative: your suite's own support predicate
    self.skipTest("cuDNN SDPA decode tier not dispatched; set ORT_TEST_REQUIRE_CUDNN_SDPA=1 to enforce")
# then assert dispatch == CUDNN_FLASH_ATTENTION unconditionally

Run both ways to prove it works AND bites:

python -m pytest <file>.py -v                              # normal: skips where unsupported
ORT_TEST_REQUIRE_CUDNN_SDPA=1 python -m pytest <file>.py -v # enforced: fails if not cuDNN

Prove the teeth. A canary you never watched fail is not verified. Force MATH-only by setting the CUDA provider's sdpa_kernel provider option to the MATH bitmask (16) — a monkeypatch of the C++ selector is not reachable from Python — under ORT_TEST_REQUIRE_CUDNN_SDPA=1, and confirm it fails with, verbatim:

AssertionError: 'CUDNN_FLASH_ATTENTION' != 'MATH'

A run that never demonstrates this failure has not proven the canary has teeth (grounding rule: negative/teeth evidence must actually be observed, not asserted).

4. Putting it together — one clean run block

WORKDIR=$(mktemp -d); cd "$WORKDIR"       # neutral + private (see §1)
source /abs/repo/.venv/bin/activate
export LD_PRELOAD=/abs/cuda12.9/lib64/libcudart.so.12:/abs/cudnn9.8/lib/libcudnn.so.9
export LD_LIBRARY_PATH=/abs/cuda12.9/lib64:/abs/cudnn9.8/lib
export PYTHONPATH=/abs/repo/onnxruntime/test/python/transformers
F=/abs/repo/onnxruntime/test/python/transformers/<file>.py

python -m pytest "$F" -v                               # A: normal
ORT_TEST_REQUIRE_CUDNN_SDPA=1 python -m pytest "$F" -v  # B: canary active, non-skippable
# C: teeth — force MATH under the env var, expect the AssertionError above

Check the passed count, not just the exit code. pytest -v exits 0 even if every test skipped (no CUDA, or an unmet @skipUnless(ml_dtypes) guard) — the saved log then looks like passing evidence but proves nothing. Require a non-zero passed count and zero unexpected skips, and note pytest exit code 5 = "no tests collected" (usually a wrong path or -k filter, not success). RUN B's canary only converts dispatch-related skips into failures — it does not rescue collection or environment skips, so still read the summary line.

Redirect to a log (... 2>&1 | tee "$WORKDIR/gpu_run.log") — the debug-info stdout and pytest output are large, and a saved log is the evidence that the run happened and dispatched to cuDNN. Write it inside $WORKDIR (the mktemp -d above), not a predictable /tmp/gpu_run.log a co-tenant could pre-create as a symlink to clobber.

Gotcha quick-reference

Symptom	Root cause	Fix
ModuleNotFoundError: onnxruntime.capi	repo-root ./onnxruntime/ source shadows the wheel	run pytest from a private mktemp -d (not repo root, not bare /tmp); abs path + PYTHONPATH
routes to MATH / cuDNN load error	torch's bundled CUDA/cuDNN shadow ORT's	LD_PRELOAD system libcudart.so.12 + libcudnn.so.9, set LD_LIBRARY_PATH
test passes but path unproven	MATH fallback gives same numbers	observe SdpaKernel= via ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO=1
all tests skip green, regression hidden	dispatch-gated skip	ORT_TEST_REQUIRE_CUDNN_SDPA=1 makes assertions non-skippable