[ericspod]

Fixes # .

Description

This adds a testing runner for the KCL GPU runner available to MONAI. The goal is to develop a GPU testing process, either as a per-PR test or nightly GPU tests.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[vikasreddy11]

Hi @ericspod , I'm interested in helping with the GPU test runner. I noticed the PR is still in progress. Is there a specific area where an independent investigation or contribution would be useful?

[ericspod]

Thanks @vikasreddy11 do you have access to a GPU? I could use some help testing the GPU resampling accuracy with the compiled extensions. I'm seeing affine tests failing like this

Details

FAIL: test_affine_003 (tests.transforms.test_affine.TestAffine.test_affine_003)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/usr/local/lib/python3.12/dist-packages/parameterized/parameterized.py", line 620, in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/monai/tests/transforms/test_affine.py", line 197, in test_affine
    test_resampler_lazy(
  File "/opt/monai/tests/lazy_transforms_utils.py", line 77, in test_resampler_lazy
    assert_allclose(lazy_out, non_lazy_out, rtol=rtol, atol=atol)
  File "/opt/monai/tests/test_utils.py", line 167, in assert_allclose
    np.testing.assert_allclose(actual, desired, *args, **kwargs)
  File "/usr/local/lib/python3.12/dist-packages/numpy/testing/_private/utils.py", line 1688, in assert_allclose
    assert_array_compare(compare, actual, desired, err_msg=str(err_msg),
  File "/usr/lib/python3.12/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/numpy/testing/_private/utils.py", line 889, in assert_array_compare
    raise AssertionError(msg)
AssertionError: 
Not equal to tolerance rtol=0.001, atol=0.001

Mismatched elements: 15 / 16 (93.8%)
Max absolute difference among violations: 1.3333333
Max relative difference among violations: 1.7777779
 ACTUAL: array([[[0.      , 0.231481, 0.601852, 0.694444],
        [0.462963, 0.833333, 1.277778, 1.296296],
        [1.203704, 1.722222, 2.166667, 2.037037],
        [1.388889, 1.898148, 2.268518, 2.083333]]], dtype=float32)
 DESIRED: array([[[0.      , 0.083333, 0.416667, 0.25    ],
        [0.166667, 0.5     , 1.166667, 0.666667],
        [0.833333, 1.833333, 2.5     , 1.333333],
        [0.5     , 1.083333, 1.416667, 0.75    ]]], dtype=float32)

This occurs in tests.transforms.test_affine and tests.transforms.test_affined. If you can try testing these files in the Docker image on in a similar environment with the CUDA extensions built and get to the bottom of the issue that would go a long way for the testing process. Be sure to have BUILD_MONAI=1 exported even when you run the tests to be sure compiled code is used. I encountered some other issues here that's related to the testing host I think that only I can sort out.

[vikasreddy11]

Thanks @ericspod, I tested on a GPU outside Docker (RTX 3050 Laptop, WSL2, Python 3.12.15, PyTorch 2.14.1+cu132, dev at cb35d45).

I can reproduce it, and it depends on the compiled extension:

• Default (USE_COMPILED=False): test_affine.py gives 125 passed, 0 failed.
• BUILD_MONAI=1 (USE_COMPILED=True): 72 failed, including test_affine_003. test_affine_116 fails alone and via unittest with the same mismatch numbers as yours.
• The failure is at the final pixel comparison in test_resampler_lazy (lazy_transforms_utils.py:77), after the pending-affine and shape assertions have passed. So the lazy metadata matches eager, and only the values resampled by apply_pending differ.

Not yet checked: test_affined, and PyTorch 2.9.x (mine is newer than CI's range). I'll report back.

[vikasreddy11]

Update: test_affined.py behaves the same way. Default (USE_COMPILED=False): 120 passed, 0 failed. BUILD_MONAI=1 (USE_COMPILED=True): 60 failed, 60 passed.

I also made a minimal standalone reproduction mirroring test_resampler_lazy (90° rotation, arange(4).reshape(1,2,2) to a 4x4 output, zero padding, align_corners=True). Eager output is identical in both modes, but lazy (via apply_pending with the test's overrides) differs only when the compiled extension is active, with the same values as the ACTUAL array in the test_affine_116 failure.

It is not GPU-specific: running the same script with CUDA_VISIBLE_DEVICES="" and BUILD_MONAI=1 gives the same wrong lazy values on CPU, so it should be reproducible without a GPU.

With align_corners=False eager and lazy match, but in lazy resample (lazy/utils.py:207) that case takes the permute/flip shortcut and never reaches SpatialResample, so it doesn't exercise the resampling.

I suspect the compiled branch of spatial_resample (spatial/functional.py ~L212) but haven't confirmed it. Still to check: PyTorch 2.9.x.

import numpy as np, torch
from monai.data import MetaTensor
from monai.transforms import Affine
from monai.transforms.lazy.functional import apply_pending
from tests.lazy_transforms_utils import get_apply_param

img = MetaTensor(torch.arange(4, dtype=torch.float32).reshape(1, 2, 2))
rot = torch.tensor([[0., -1., 0.], [1., 0., 0.], [0., 0., 1.]])
init = dict(affine=rot, padding_mode="zeros", align_corners=True)
call = dict(img=img, spatial_size=(4, 4))
eager = Affine(**init)(**call)
eager = eager[0] if isinstance(eager, tuple) else eager
t = Affine(**init); t.lazy = True
out = t(**call)
out = out[0] if isinstance(out, tuple) else out
out = apply_pending(out, overrides=get_apply_param(init, call))[0]
print(np.round(np.asarray(eager), 4)); print(np.round(np.asarray(out), 4))

Run with PYTHONPATH=. from the repo root, plain and with BUILD_MONAI=1.

I'd like to dig into this. If I find a fix, would you like me to open a PR, or is someone already working on this?

[ericspod]

Hi @vikasreddy11 thanks for looking into the issue. If you could please find a fix and propose a PR if you do, no one else is working on it currently. These fails are the major issues we're encountering with the GPU tests but I'm working on others currently. The root cause hasn't been found since the PR tests compile the code but then don't use it since BUILD_MONAI=1 must be present even at runtime, this behaviour is something I'll want to change here.

[ericspod]

The issue with the GlobalMutualInformationLoss tests is that torch.bmm tries to compile with Triton when used with 1D matrices and fails because the container for the action doesn't have a compiler. This example raises the issue in PyTorch 2.14.1:

import torch
t = torch.rand(1, 10, 1).cuda()
tt = torch.rand(1, 1, 10).cuda()
torch.bmm(t, tt)

This raises "RuntimeError: Failed to find C compiler." on line 364 of image_dissimilarity.py. I would expect PyTorch to fall back to a native computation, the solution here is to move the inputs to CPU then move the result back to CUDA. Another solution is a native matmul function just for this loss:

t = torch.rand(1, 10, 1).cuda()
tt = torch.rand(1, 1, 10).cuda()

def matmul_native(a, b):  #(B,N,1) (B,1,M)
    return a.repeat(1, 1, b.shape[2]) * b.repeat(1, a.shape[1], 1)

matmul_native(t,tt)