Skip to content

fix(tests): compare baddbmm lowering against TensorRT in fp32, not TF32 - #4586

Open
shoumikhin wants to merge 2 commits into
pytorch:mainfrom
shoumikhin:fix/baddbmm-tf32
Open

fix(tests): compare baddbmm lowering against TensorRT in fp32, not TF32#4586
shoumikhin wants to merge 2 commits into
pytorch:mainfrom
shoumikhin:fix/baddbmm-tf32

Conversation

@shoumikhin

Copy link
Copy Markdown
Contributor

Problem

dynamo-lowering-rtx is red on main:

FAILED lowering/test_decompositions.py::TestLowering::test_lowering_baddbmm
AssertionError: 0.0019378662109375 != 0 within 4 places
: baddbmm TRT outputs don't match with the original model.

The same test passes on standard TensorRT, so it went in green and only turned red once the RTX lane ran on main.

Cause

The test compares a TensorRT engine against an eager fp32 reference to DECIMALS_OF_AGREEMENT, but builds the engine with TF32 left at its default, which is on. TF32 rounds both matmul operands to 10 mantissa bits before multiplying. A passing run therefore only means TensorRT happened to pick a kernel that was not TF32; it is not something the test asks for. TensorRT-RTX picks a TF32 kernel and the check fails.

The decomposition is not at fault. Over 500 random draws of the shapes the test uses, bias + bmm(batch1, batch2) in fp32 matches aten.baddbmm bit for bit (max difference 0.0). Repeating the same math with the two operands rounded to 10 mantissa bits drifts by up to 0.0054, which brackets the 0.0019 seen in CI.

Fix

Pass disable_tf32=True so both sides run the same arithmetic. The tolerance is unchanged and nothing is skipped. This matches how other numeric tests in the repo handle it, for example tests/py/dynamo/conversion/test_index_put_aten.py and tests/py/dynamo/models/test_models.py.

Testing

test_lowering_baddbmm passes locally on standard TensorRT with the change, and the debug log confirms disable_tf32=True reaches the compilation settings through the torch_compile path. The RTX lane needs ci: full or an approving review to run here.

@meta-cla meta-cla Bot added the cla signed label Aug 26, 2026
@github-actions github-actions Bot added the component: tests Issues re: Tests label Aug 26, 2026
@github-actions
github-actions Bot requested a review from lanluo-nvidia August 26, 2026 02:49
test_lowering_baddbmm builds the engine with TF32 left at its default (on),
then checks the result against an eager fp32 reference to DECIMALS_OF_AGREEMENT.
TF32 rounds the matmul operands to 10 mantissa bits, so a passing run only means
TensorRT happened to pick a non-TF32 kernel. TensorRT-RTX picks a TF32 one and
the test fails on main with a difference of 0.0019378662109375.

The decomposition itself is exact: over 500 random draws of the shapes the test
uses, bias + bmm(batch1, batch2) in fp32 matches aten.baddbmm bit for bit, while
the same math with the operands rounded to TF32 drifts up to 0.0054.

Ask TensorRT for fp32 so both sides run the same arithmetic. The tolerance and
the coverage stay as they are.
The repository lint job runs `black --check .` across the whole tree, so any
file that does not match the formatter fails CI for every open pull request,
not only the one that touched it. `tests/py/dynamo/conversion/test_cumsum_aten.py`
is currently not black-conformant on main, which turns the Python Linting check
red here.

Reformat that one file with black. This is a formatting-only change: two
statements that fit on a single line are un-wrapped. No test logic changes.

Verified by running `black --check .` on the full tree: all files pass.
@lanluo-nvidia lanluo-nvidia added this to the v2.14.0 milestone Aug 26, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants