fix(asr): repair ConvSubsampling forward paths missed by the MaskedConvSequential refactor

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

@@ -492,7 +492,13 @@ def forward(self, x, lengths):
             else:
                 x, lengths = self.conv(x, lengths)
         else:
-            x, lengths = self.conv(x)
+            if self.conv2d_subsampling:
+                # Chunking disabled (-1): run the masked conv stack directly.
+                x, lengths = self.conv(x, lengths)
+            else:
+                # 1-D conv stacks run without masking; keep the lengths computed above.
+                x = self.conv(x)
+                lengths = out_lengths
 
         # Flatten Channel and Frequency Axes
         if self.conv2d_subsampling:
@@ -725,7 +731,15 @@ class MaskedConvSequential(nn.Sequential):
     # Set by ConvSubsampling; off by default, so every other subsampling type stays on PyTorch.
     fuse_triton = False
 
-    def forward(self, x, lengths):
+    def forward(self, x, lengths=None):
+        if lengths is None:
+            # Plain pass-through, equivalent to nn.Sequential. Used by the 1-D conv
+            # stacks (striding_conv1d / dw_striding_conv1d), whose inputs are already
+            # channels-first and which do not take part in length masking.
+            for layer in self:
+                x = layer(x)
+            return x
+
         # Convert input (batch, time, features) to conv format
         x = x.unsqueeze(1)  # (batch, 1, time, features)
         current_lengths = lengths
@@ -760,10 +774,22 @@ def _forward_torch(self, x, current_lengths):
             x = layer(x)
 
             # Update lengths for stride operations with proper padding
-            if hasattr(layer, 'stride') and layer.stride != (1, 1):
-                current_lengths = calculate_conv_output_size(
-                    current_lengths, layer.kernel_size[0], layer.stride[0], _layer_padding(layer)
-                )
+            if hasattr(layer, 'stride') and _pair_first(layer.stride) != 1:
+                kernel_size = _pair_first(layer.kernel_size)
+                stride = _pair_first(layer.stride)
+                left_pad, right_pad = _layer_padding(layer)
+                if getattr(layer, 'ceil_mode', False):
+                    # Ceil-mode pooling (e.g. vggnet's MaxPool2d) emits one extra frame
+                    # whenever the floor division drops a remainder.
+                    remainder = (current_lengths + left_pad + right_pad - kernel_size) % stride
+                    current_lengths = calculate_conv_output_size(
+                        current_lengths, kernel_size, stride, (left_pad, right_pad)
+                    )
+                    current_lengths = current_lengths + (remainder != 0).long()
+                else:
+                    current_lengths = calculate_conv_output_size(
+                        current_lengths, kernel_size, stride, (left_pad, right_pad)
+                    )
                 mask = self._create_mask(x, current_lengths.long())
 
         return x, current_lengths, mask
@@ -826,15 +852,26 @@ def _create_mask(self, tensor, lengths):
         return time_mask.unsqueeze(-1).expand(batch_size, time, features).to(tensor.dtype)
 
 
+def _pair_first(value):
+    """First element of an int-or-tuple kernel/stride/padding attribute.
+
+    nn.Conv2d stores tuples, but the pooling modules (nn.MaxPool2d/nn.AvgPool2d) keep
+    whatever was passed in, which ConvSubsampling passes as a plain int.
+    """
+    return value[0] if isinstance(value, tuple) else value
+
+
 def _layer_padding(layer):
     """The (start, end) padding of a convolution.
 
-    nn.Conv2d's `.padding` is (pad_h, pad_w), one value per axis and symmetric within it, so the
-    height value is both edges. CausalConv2D keeps its two edges on private attributes.
+    nn.Conv2d's `.padding` is (pad_h, pad_w), one value per axis and symmetric within it, so
+    the height value is both edges. CausalConv2D keeps its two edges on private attributes.
+    Pooling layers store an int (or an int-per-axis tuple), which is symmetric too.
     """
     if hasattr(layer, "_left_padding"):
         return layer._left_padding, layer._right_padding
-    return layer.padding[0], layer.padding[0]
+    padding = _pair_first(layer.padding)
+    return padding, padding
 
 
 def _is_depthwise(layer):

tests/collections/asr/test_asr_subsampling.py

@@ -16,6 +16,7 @@
 import torch
 
 from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.parts.submodules.subsampling import ConvSubsampling, calc_length
 
 
 class TestASRSubsamplingConvChunking:
@@ -60,3 +61,106 @@ def test_forward(self):
         assert diff <= 0.2
         diff = torch.mean(torch.abs(logprobs_batch4_split - logprobs_batch4_nosplit))
         assert diff <= 0.2
+
+
+class TestConvSubsamplingForwardPaths:
+    """CPU tests for ConvSubsampling paths that the chunking/splitting tests do not cover.
+
+    Covers `subsampling_conv_chunking_factor=-1` (chunking disabled), the 1-D conv
+    stacks, and ceil-mode length bookkeeping for vggnet pooling.
+    """
+
+    @pytest.mark.run_only_on('CPU')
+    @pytest.mark.unit
+    @pytest.mark.parametrize("subsampling", ["striding", "dw_striding", "vggnet"])
+    @pytest.mark.parametrize("subsampling_factor", [4, 8])
+    def test_no_chunking_forward_matches_chunked_path(self, subsampling, subsampling_factor):
+        """With chunking disabled (-1), forward must run and match the default path."""
+        torch.manual_seed(0)
+        no_chunk = ConvSubsampling(
+            subsampling=subsampling,
+            subsampling_factor=subsampling_factor,
+            feat_in=80,
+            feat_out=64,
+            conv_channels=32,
+            subsampling_conv_chunking_factor=-1,
+        ).eval()
+        default = ConvSubsampling(
+            subsampling=subsampling,
+            subsampling_factor=subsampling_factor,
+            feat_in=80,
+            feat_out=64,
+            conv_channels=32,
+        ).eval()
+        default.load_state_dict(no_chunk.state_dict())
+
+        x = torch.randn(2, 101, 80)
+        lengths = torch.tensor([101, 97])
+
+        with torch.inference_mode():
+            out, out_lengths = no_chunk(x, lengths)
+            ref_out, ref_lengths = default(x, lengths)
+
+        assert out.shape == ref_out.shape
+        assert torch.allclose(out, ref_out)
+        assert out_lengths.tolist() == ref_lengths.tolist()
+        assert bool((out_lengths <= out.shape[1]).all())
+
+    @pytest.mark.run_only_on('CPU')
+    @pytest.mark.unit
+    @pytest.mark.parametrize("subsampling", ["striding_conv1d", "dw_striding_conv1d"])
+    @pytest.mark.parametrize("subsampling_factor", [2, 4, 8])
+    def test_conv1d_stacks_forward(self, subsampling, subsampling_factor):
+        """The 1-D conv stacks must run with the default config and report calc_length lengths."""
+        module = ConvSubsampling(
+            subsampling=subsampling,
+            subsampling_factor=subsampling_factor,
+            feat_in=80,
+            feat_out=64,
+            conv_channels=32,
+        ).eval()
+
+        x = torch.randn(2, 101, 80)
+        lengths = torch.tensor([101, 97])
+
+        with torch.inference_mode():
+            out, out_lengths = module(x, lengths)
+
+        sampling_num = module._sampling_num
+        expected = calc_length(
+            lengths=lengths.to(dtype=torch.float),
+            all_paddings=module._left_padding + module._right_padding,
+            kernel_size=module._kernel_size,
+            stride=module._stride,
+            ceil_mode=module._ceil_mode,
+            repeat_num=sampling_num,
+        )
+        assert out.shape[0] == 2
+        assert out.shape[2] == 64
+        assert out_lengths.tolist() == expected.tolist()
+
+    @pytest.mark.run_only_on('CPU')
+    @pytest.mark.unit
+    @pytest.mark.parametrize("input_length", [101, 97])
+    def test_vggnet_lengths_use_ceil_mode(self, input_length):
+        """vggnet pooling runs with ceil_mode=True; reported lengths must match it."""
+        module = ConvSubsampling(
+            subsampling='vggnet', subsampling_factor=4, feat_in=80, feat_out=64, conv_channels=32
+        ).eval()
+
+        x = torch.randn(1, input_length, 80)
+        lengths = torch.tensor([input_length])
+
+        with torch.inference_mode():
+            out, out_lengths = module(x, lengths)
+
+        expected = calc_length(
+            lengths=lengths.to(dtype=torch.float),
+            all_paddings=0,
+            kernel_size=2,
+            stride=2,
+            ceil_mode=True,
+            repeat_num=2,
+        )
+        assert out.shape[1] == expected.item()
+        assert out_lengths.tolist() == expected.tolist()