diff --git a/nemo/collections/asr/parts/submodules/subsampling.py b/nemo/collections/asr/parts/submodules/subsampling.py index d3a21aae53c4..905383cc0821 100644 --- a/nemo/collections/asr/parts/submodules/subsampling.py +++ b/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) diff --git a/tests/collections/asr/test_asr_subsampling.py b/tests/collections/asr/test_asr_subsampling.py index 8f638afb73c0..c561dac4d350 100644 --- a/tests/collections/asr/test_asr_subsampling.py +++ b/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