Fix 32-bit indexing overflow in conv subsampling chunking guard

nemo/collections/asr/parts/submodules/subsampling.py

@@ -101,6 +101,13 @@ def forward(self, x, lengths):
         return x, lengths
 
 
+# cuDNN and PyTorch's native CUDA kernels index tensor elements with 32-bit integers, so
+# any tensor entering or leaving a conv must hold fewer than this many elements; exceeding
+# it raises "Expected canUse32BitIndexMath(...) to be true, but got false".
+# See https://github.com/pytorch/pytorch/issues/80020
+_MAX_CONV_NUMEL_32BIT = 2**31 - 1
+
+
 class ConvSubsampling(torch.nn.Module):
     """Convolutional subsampling which supports VGGNet and striding approach introduced in:
     VGGNet Subsampling: Transformer-transducer: end-to-end speech recognition with self-attention (https://arxiv.org/pdf/1910.12977.pdf)
@@ -424,6 +431,22 @@ def get_sampling_frames(self):
     def get_streaming_cache_size(self):
         return [0, self.subsampling_factor + 1]
 
+    def _first_conv_output_numel(self, x):
+        """Elements in the first conv layer's output - the largest activation in the stack.
+
+        The first conv (1 -> conv_channels, strided in both T and F) produces the biggest
+        tensor for the 'striding' and 'dw_striding' variants; every later layer only shrinks
+        T and F. ``x`` is the ``(B, T, F)`` input before the channel dim is added.
+
+        Note: this assumes a strided first conv. The 'vgg' variant starts with stride-1
+        convs, so its largest activation is not bounded by this estimate.
+        """
+        b, t, f = x.size()
+        pad = (self._left_padding, self._right_padding)
+        out_t = calculate_conv_output_size(t, self._kernel_size, self._stride, pad)
+        out_f = calculate_conv_output_size(f, self._kernel_size, self._stride, pad)
+        return b * self._conv_channels * out_t * out_f
+
     def forward(self, x, lengths):
         out_lengths = calc_length(
             lengths,
@@ -444,11 +467,14 @@ def forward(self, x, lengths):
                 # if subsampling_conv_chunking_factor is 1, we split only if needed
                 # avoiding a bug / feature limiting indexing of tensors to 2**31
                 # see https://github.com/pytorch/pytorch/issues/80020
-                x_ceil = 2**31 / self._conv_channels * self._stride * self._stride
-                if torch.numel(x) > x_ceil:
-                    need_to_split = True
-                else:
-                    need_to_split = False
+                # Compare the exact first-conv output (the largest activation in the
+                # stack) against the hard 32-bit element limit, splitting on '>=': at
+                # equality the tensor already has INT_MAX elements, which is exactly what
+                # trips canUse32BitIndexMath. The conv input is guarded too.
+                need_to_split = (
+                    self._first_conv_output_numel(x) >= _MAX_CONV_NUMEL_32BIT
+                    or torch.numel(x) >= _MAX_CONV_NUMEL_32BIT
+                )
             else:
                 # if subsampling_conv_chunking_factor > 1 we always split
                 need_to_split = True
@@ -512,9 +538,12 @@ def conv_split_by_batch(self, x, lengths):
         else:
             # avoiding a bug / feature limiting indexing of tensors to 2**31
             # see https://github.com/pytorch/pytorch/issues/80020
-            x_ceil = 2**31 / self._conv_channels * self._stride * self._stride
-            p = math.ceil(math.log(torch.numel(x) / x_ceil, 2))
-            cf = 2**p
+            # Smallest power-of-two batch split that keeps each chunk strictly below the
+            # 32-bit element limit (the +1 forces "strictly"). Mirror the forward() guard
+            # by sizing against whichever of the conv input or its (larger) first-conv
+            # output reaches the limit.
+            numel = max(self._first_conv_output_numel(x), torch.numel(x))
+            cf = 2 ** math.ceil(math.log(numel // _MAX_CONV_NUMEL_32BIT + 1, 2))
             logging.debug(f'using auto set chunking factor: {cf}')
 
         new_batch_size = b // cf
@@ -549,8 +578,9 @@ def conv_split_by_channel(self, x):
             else:
                 # avoiding a bug / feature limiting indexing of tensors to 2**31
                 # see https://github.com/pytorch/pytorch/issues/80020
-                p = math.ceil(math.log(torch.numel(x) / 2**31, 2))
-                cf = 2**p
+                # +1 keeps each chunk strictly below the 32-bit element limit and avoids a
+                # fractional factor when the tensor is already within the limit.
+                cf = 2 ** math.ceil(math.log(torch.numel(x) // _MAX_CONV_NUMEL_32BIT + 1, 2))
                 logging.debug(f'using auto set chunking factor: {cf}')
 
             new_c = int(c // cf)

tests/collections/asr/test_asr_subsampling.py

@@ -11,10 +11,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import math
+
 import pytest
 import torch
 
 from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.parts.submodules import subsampling as subsampling_module
+from nemo.collections.asr.parts.submodules.subsampling import ConvSubsampling
 
 
 class TestASRSubsamplingConvChunking:
@@ -59,3 +63,109 @@ def test_forward(self):
         assert diff <= 0.2
         diff = torch.mean(torch.abs(logprobs_batch4_split - logprobs_batch4_nosplit))
         assert diff <= 0.2
+
+
+def _build_conv_subsampling(feat_in=16, conv_channels=8, factor=4):
+    """A tiny dw_striding ConvSubsampling for unit-testing the 32-bit chunking logic."""
+    return ConvSubsampling(
+        subsampling="dw_striding",
+        subsampling_factor=factor,
+        feat_in=feat_in,
+        feat_out=32,
+        conv_channels=conv_channels,
+        subsampling_conv_chunking_factor=1,
+    ).eval()
+
+
+def _install_split_spy(monkeypatch, sub):
+    """Record the batch size of every conv_split_by_batch call; returns the list."""
+    calls = []
+    original_split = sub.conv_split_by_batch
+
+    def spy_split(inp, lens):
+        calls.append(int(inp.shape[0]))
+        return original_split(inp, lens)
+
+    monkeypatch.setattr(sub, "conv_split_by_batch", spy_split)
+    return calls
+
+
+class TestConvSubsampling32BitIndexing:
+    """Unit tests for the exact 32-bit element-limit guard and auto chunking factor.
+
+    These run on small synthetic inputs with the limit lowered via monkeypatch, so they
+    exercise the splitting logic without allocating multi-GB tensors.
+    """
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("shape", [(1, 7, 16), (3, 50, 16), (5, 123, 16)])
+    def test_first_conv_output_numel_matches_real_conv(self, shape):
+        # The estimate must equal the actual element count of the first conv's output,
+        # which is the largest activation the 32-bit limit is checked against.
+        sub = _build_conv_subsampling()
+        x = torch.randn(*shape)
+        real_numel = sub.conv[0](x.unsqueeze(1)).numel()  # run only the first Conv2d
+        assert sub._first_conv_output_numel(x) == real_numel
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("batch_size", [1, 4])
+    def test_guard_splits_at_exact_limit(self, monkeypatch, batch_size):
+        # At output == limit the split must trigger. The previous '>' guard let a tensor of
+        # exactly INT_MAX elements (the value that trips canUse32BitIndexMath) through unsplit.
+        sub = _build_conv_subsampling()
+        x = torch.randn(batch_size, 50, 16)
+        lengths = torch.full((batch_size,), 50, dtype=torch.long)
+
+        # Reference with the real (large) limit: no splitting happens.
+        ref, ref_len = sub(x.clone(), lengths.clone())
+
+        split_calls = _install_split_spy(monkeypatch, sub)
+        monkeypatch.setattr(subsampling_module, "_MAX_CONV_NUMEL_32BIT", sub._first_conv_output_numel(x))
+
+        out, out_len = sub(x.clone(), lengths.clone())
+
+        assert split_calls, "the guard did not split when the first-conv output equals the 32-bit limit"
+        # Splitting (by batch, or by channel when batch_size == 1) must not change the result.
+        assert torch.allclose(out, ref, atol=1e-5)
+        assert torch.equal(out_len, ref_len)
+
+    @pytest.mark.unit
+    def test_guard_does_not_split_below_limit(self, monkeypatch):
+        # One element below the limit must not split: no needless chunking.
+        sub = _build_conv_subsampling()
+        x = torch.randn(4, 50, 16)
+        lengths = torch.full((4,), 50, dtype=torch.long)
+
+        split_calls = _install_split_spy(monkeypatch, sub)
+        monkeypatch.setattr(subsampling_module, "_MAX_CONV_NUMEL_32BIT", sub._first_conv_output_numel(x) + 1)
+
+        sub(x.clone(), lengths.clone())
+        assert not split_calls
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("batch_size", [4, 8, 16])
+    def test_auto_chunking_keeps_each_chunk_below_limit(self, monkeypatch, batch_size):
+        # The auto chunking factor must split into chunks whose first-conv output is strictly
+        # below the limit (the previous float formula could pick a fractional factor or leave a
+        # chunk sitting exactly at the limit).
+        sub = _build_conv_subsampling()
+        x = torch.randn(batch_size, 40, 16)
+        lengths = torch.full((batch_size,), 40, dtype=torch.long)
+        limit = sub._first_conv_output_numel(x) // 3 + 1  # forces a multi-way split
+        monkeypatch.setattr(subsampling_module, "_MAX_CONV_NUMEL_32BIT", limit)
+
+        # Record the batch size of each chunk actually fed to the conv stack.
+        chunk_batches = []
+        original_forward = sub.conv.forward
+
+        def recording_forward(inp, lens):
+            chunk_batches.append(int(inp.shape[0]))
+            return original_forward(inp, lens)
+
+        monkeypatch.setattr(sub.conv, "forward", recording_forward)
+
+        sub(x.clone(), lengths.clone())
+
+        assert len(chunk_batches) > 1, "expected the input to be split into multiple chunks"
+        for chunk_batch in chunk_batches:
+            assert sub._first_conv_output_numel(x[:chunk_batch]) < limit