Fix bus error or segfault from roi_align with large batchsize by zy1git · Pull Request #9441 · pytorch/vision

zy1git · 2026-03-13T09:51:18Z

Summary
Bug: roi_align in torchvision crashes with a bus error/segfault on CPU or returns silently wrong (all-zero) results on CUDA when the total number of output elements exceeds INT_MAX (~2.1 billion). This is caused by 32-bit int overflow in index arithmetic within the C++ and CUDA kernels.

Root Cause: The kernels use int for composite index calculations like n × channels × pooled_width × pooled_height and pointer offsets like (roi_batch_ind × channels + c) × height × width. When these products exceed 2,147,483,647, the int wraps to a negative value, causing out-of-bounds memory access.

Example: FasterRCNN with batch_size=172 generates ~172,000 ROIs. The output index reaches 171,999 × 256 × 7 × 7 = 2,157,555,456 > INT_MAX, which matches the reporter's observed threshold exactly.

Fix: Promoted int to int64_t for all index, offset, and stride variables in the relevant files.

Test Plan
New overflow regression test
pytest test/test_ops.py::TestRoIAlign::test_roi_align_large_index -v

Existing tests — verify no regressions
pytest test/test_ops.py::TestRoIAlign -v

Fixes #8206

pytorch-bot · 2026-03-13T09:51:23Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9441

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit b0004a1 with merge base 6285457 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

NicolasHug · 2026-03-16T15:46:56Z

test/test_ops.py

+        output_bytes = n_rois * channels * pooled_h * pooled_w * 4  # float32
+        if output_bytes > 9 * 1024**3:
+            pytest.skip("Test requires ~9 GB of memory")
+


all these values are statically defined. This if block is either always True or always False.

Thanks for pointing this out. I agree. I removed that part in the new commit.

NicolasHug · 2026-03-16T15:47:28Z

test/test_ops.py

+            x = torch.rand(num_imgs, channels, height, width, dtype=torch.float32, device=device)
+            rois = torch.zeros(n_rois, 5, dtype=torch.float32, device=device)
+        except RuntimeError:
+            pytest.skip("Not enough memory to allocate test tensors")


Please verify tests aren't not being skipped on the CI. If they pass, remove the try/except, if they don't, we'll have to consider other strategies to test this.

Thanks for the comment. I verified that the tests are not being skipped on the CI and my devserver. I removed the try/except in the new commit.

NicolasHug · 2026-03-16T15:49:45Z

torchvision/csrc/ops/cpu/roi_align_kernel.cpp

 template <typename T>
 void roi_align_backward_kernel_impl(
-    int nthreads,
+    int64_t nthreads,


Can you explain why nthreads needs to be int64_t? It should never need to be that large? If it's for integer comparison to not warn, we could just cast?

Thanks for the question. nthreads in the backward kernel is grad.numel(), which equals n_rois × channels × pooled_h × pooled_w.

I changed to int nthreads in both .cpp and .cu files and ran the test (only the forward kernel test, no backward kernel test due to the large memory requirement). The CPU test passed but the CUDA test failed with all-zero output. The reason is that the CPU forward kernel doesn't use nthreads — it loops over n_rois separately, and the overflow is handled by the int64_t changes to index_n, index_n_c, and index. The CUDA forward kernel uses a flat loop with nthreads as the bound, so truncating to int caused nthreads to wrap to a negative value, making the loop condition immediately false and skipping all output computation — resulting in all-zero output.

The CPU backward kernel does use nthreads in the same flat-loop pattern as CUDA (for (int64_t index = 0; index < nthreads; ...)) and receives the same overflowing value via grad.numel(), so it needs int64_t for the same reason.

A backward-specific test would require large memory (output + grad_output + grad_input), which might be impractical for CI. Do we want to add one with a memory skip guard, or is the current forward-only test sufficient?

As for "If it's for integer comparison to not warn, we could just cast?", I think it isn't a warning issue. The problem is that the value could be actually large and gets truncated at the call site before the function body runs. In the author's reproducing example (batch_size=172, default 1000 proposals per image), nthreads is 172,000 × 256 × 7 × 7 = 2,157,568,000 > INT_MAX.

I added the backward kernel test in the latest commit. If I change int64_t nthreadsto int nthreads, the CPU backward test fails with all-zero gradients because nthreads gets truncated to a negative value and the loop never executes.

nthreads doesn't mean "number of threads" in this code, instead, it means "total number of output elements to process.", which could be very large.

Feel free to let me know if you have any questions.

Thanks. For my own ref:

nthreads is the name for output_size in the cuda kernel:

vision/torchvision/csrc/ops/cuda/roi_align_kernel.cu

Lines 378 to 379 in d7400a3

roi_align_forward_kernel_impl<scalar_t><<<grid, block, 0, stream>>>(

output_size,

where output size was already inferred as uint64 since size() returns uint64:

vision/torchvision/csrc/ops/cuda/roi_align_kernel.cu

Lines 354 to 362 in d7400a3

auto num_rois = rois.size(0);

auto channels = input.size(1);

auto height = input.size(2);

auto width = input.size(3);

at::Tensor output = at::zeros(

{num_rois, channels, pooled_height, pooled_width}, input.options());

auto output_size = num_rois * pooled_height * pooled_width * channels;

nthreads was previously implicitly cast to int32 when passed to roi_align_forward_kernel_impl causing an overflow. Now, with the change made to the function signature, it's properly kept as int64.

NicolasHug · 2026-03-23T09:42:43Z

test/test_ops.py


+    @pytest.mark.parametrize("device", cpu_and_cuda())
+    def test_roi_align_large_index(self, device):
+        """Regression test for https://github.com/pytorch/vision/issues/8206"""


Suggested change

"""Regression test for https://github.com/pytorch/vision/issues/8206"""

"""Non-regression test for https://github.com/pytorch/vision/issues/8206"""

NicolasHug · 2026-03-23T09:44:38Z

test/test_ops.py

+
+        # Forward kernel test
+        assert result.shape == (n_rois, channels, pooled_h, pooled_w)
+        assert result.abs().sum() > 0, "roi_align returned all zeros — likely an index overflow bug"


Here and below, don't specify anything beyond the assert, pytest is already good at showing the right thing. Here, the message doesn't help that much either.

Suggested change

assert result.abs().sum() > 0, "roi_align returned all zeros — likely an index overflow bug"

assert result.abs().sum() > 0

NicolasHug · 2026-03-23T10:19:35Z

torchvision/csrc/ops/cpu/roi_align_common.h

+          pc.pos1 = static_cast<int64_t>(y_low) * width + x_low;
+          pc.pos2 = static_cast<int64_t>(y_low) * width + x_high;
+          pc.pos3 = static_cast<int64_t>(y_high) * width + x_low;
+          pc.pos4 = static_cast<int64_t>(y_high) * width + x_high;


Can you double check that these casts are needed?

claude says: y_low and y_high are pixel coordinates bounded by height, and width is the image width. y_low * width + x_low is at most height * width, which is the number of pixels in a single channel of a single image. That's not going to overflow int.

Please check every single other change in this PR.

NicolasHug · 2026-03-23T10:23:15Z

torchvision/csrc/ops/cpu/roi_align_kernel.cpp

 template <typename T>
 void roi_align_backward_kernel_impl(
-    int nthreads,
+    int64_t nthreads,


Thanks. For my own ref:

nthreads is the name for output_size in the cuda kernel:

vision/torchvision/csrc/ops/cuda/roi_align_kernel.cu

Lines 378 to 379 in d7400a3

roi_align_forward_kernel_impl<scalar_t><<<grid, block, 0, stream>>>(

output_size,

where output size was already inferred as uint64 since size() returns uint64:

vision/torchvision/csrc/ops/cuda/roi_align_kernel.cu

Lines 354 to 362 in d7400a3

auto num_rois = rois.size(0);

auto channels = input.size(1);

auto height = input.size(2);

auto width = input.size(3);

at::Tensor output = at::zeros(

{num_rois, channels, pooled_height, pooled_width}, input.options());

auto output_size = num_rois * pooled_height * pooled_width * channels;

nthreads was previously implicitly cast to int32 when passed to roi_align_forward_kernel_impl causing an overflow. Now, with the change made to the function signature, it's properly kept as int64.

NicolasHug · 2026-03-23T10:24:33Z

test/test_ops.py

        ), f"Expected execution to take < {execution_time_ms_threshold} ms, actually took {execution_time_ms} ms"

+    @pytest.mark.parametrize("device", cpu_and_cuda())
+    def test_roi_align_large_index(self, device):


This is actually taking a super long time on CPU:

https://productionresultssa18.blob.core.windows.net/actions-results/5a8fef95-8252-44a2-86fb-264662c987f6/workflow-job-run-05f658de-4299-5615-a958-010a82bf0999/logs/job/job-logs.txt?rsct=text%2Fplain&se=2026-03-23T10%3A33%3A43Z&sig=9bWAdM4OgZyShqk2hZAKOacge4p93mIQqKcfYlkUHVI%3D&ske=2026-03-23T12%3A13%3A31Z&skoid=ca7593d4-ee42-46cd-af88-8b886a2f84eb&sks=b&skt=2026-03-23T08%3A13%3A31Z&sktid=398a6654-997b-47e9-b12b-9515b896b4de&skv=2025-11-05&sp=r&spr=https&sr=b&st=2026-03-23T10%3A23%3A38Z&sv=2025-11-05

2026-03-20T08:49:14.6040945Z 260.23s call test/test_ops.py::TestRoIAlign::test_roi_align_large_index[cpu]

on GPU it seems OK. Let's just skip the test on CPU indicating it takes too long.

Zhitao Yu added 2 commits March 13, 2026 02:38

fix the issue 8206 and add the test

8c71ea8

fix the issue 8206 and add the test

40b2276

meta-cla bot added the cla signed label Mar 13, 2026

remove unnecessary comments

d9ab5ce

zy1git marked this pull request as draft March 13, 2026 10:03

NicolasHug reviewed Mar 16, 2026

View reviewed changes

Zhitao Yu and others added 4 commits March 17, 2026 02:11

address the comments

544d960

add the backward kernel test

76f6a16

address the test failure

b6b7ab7

Merge branch 'main' into issue-8206

b0004a1

zy1git marked this pull request as ready for review March 20, 2026 08:27

NicolasHug reviewed Mar 23, 2026

View reviewed changes

	roi_align_forward_kernel_impl<scalar_t><<<grid, block, 0, stream>>>(
	output_size,

	auto num_rois = rois.size(0);
	auto channels = input.size(1);
	auto height = input.size(2);
	auto width = input.size(3);

	at::Tensor output = at::zeros(
	{num_rois, channels, pooled_height, pooled_width}, input.options());

	auto output_size = num_rois * pooled_height * pooled_width * channels;

	"""Regression test for https://github.com/pytorch/vision/issues/8206"""
	"""Non-regression test for https://github.com/pytorch/vision/issues/8206"""

	assert result.abs().sum() > 0, "roi_align returned all zeros — likely an index overflow bug"
	assert result.abs().sum() > 0

Conversation

zy1git commented Mar 13, 2026

Uh oh!

pytorch-bot bot commented Mar 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9441

✅ No Failures

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NicolasHug Mar 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NicolasHug Mar 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zy1git Mar 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zy1git Mar 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pytorch-bot bot commented Mar 13, 2026 •

edited

Loading

NicolasHug Mar 16, 2026 •

edited

Loading

NicolasHug Mar 16, 2026 •

edited

Loading

zy1git Mar 17, 2026 •

edited

Loading

zy1git Mar 17, 2026 •

edited

Loading