Simplify packed PEE production controls

Author: pzelasko

nemo/collections/asr/modules/ggemm_transformer_encoder.py

@@ -512,7 +512,6 @@ def forward_grouped_sequence_packed(
         backend: str = 'baddbmm',
         moe_mode: str = 'dense',
         fused_qkv: bool = False,
-        strict: bool = False,
     ) -> Dict[str, PackedEncoderOutput]:
         """Run all experts in layer lockstep using native THD grouped kernels.
 
@@ -528,7 +527,6 @@ def forward_grouped_sequence_packed(
             backend=backend,
             moe_mode=moe_mode,
             fused_qkv=fused_qkv,
-            strict=strict,
         )
 
     def forward_grouped(
@@ -1324,7 +1322,6 @@ def _forward_grouped_sequence_packed(
         backend: str,
         moe_mode: str,
         fused_qkv: bool,
-        strict: bool,
     ) -> Dict[str, PackedEncoderOutput]:
         if backend not in GROUPED_GEMM_BACKENDS:
             raise ValueError(f"Unknown grouped-GEMM backend '{backend}'; expected one of {GROUPED_GEMM_BACKENDS}.")
@@ -1394,24 +1391,6 @@ def _forward_grouped_sequence_packed(
             }
         del prepared, padded, packed, output_lengths, shared_mask
 
-        if strict:
-            incompatibilities = []
-            if any(expert.self_attention_model == 'rel_pos' for expert in encs.values()):
-                incompatibilities.append('relative-position attention')
-            if len({expert.d_model // expert.n_heads for expert in encs.values()}) != 1:
-                incompatibilities.append('head dimensions')
-            if len({expert.attn_mode for expert in encs.values()}) != 1:
-                incompatibilities.append('attention modes')
-            if len({state[name]['metadata_key'] for name in self.expert_names}) != 1:
-                incompatibilities.append('packed sequence boundaries')
-            if len({(state[name]['x'].device, state[name]['x'].dtype) for name in self.expert_names}) != 1:
-                incompatibilities.append('devices/dtypes')
-            if incompatibilities:
-                raise ValueError(
-                    "Strict grouped sequence-packed execution requires one compatible attention bucket; "
-                    f"incompatible {', '.join(incompatibilities)}."
-                )
-
         trace = {
             'mode': 'grouped_thd',
             'layers': n_layers,

nemo/collections/asr/modules/parallel_expert_encoder.py

@@ -883,25 +883,21 @@ def forward_sequence_packed(
         length,
         spk_targets=None,
         return_experts: bool = False,
-        *,
-        grouped: bool = True,
     ) -> PackedEncoderOutput | tuple[PackedEncoderOutput, dict[str, object]]:
         """Encode offline while keeping expert Transformer activations token-flat.
 
         Online/windowed inference retains its established prefix/cache path. Existing
         :meth:`forward` remains the Conformer-compatible padded API.
 
         Set ``return_experts=True`` to also return packed speech/sound states and
-        the padded Sortformer speaker predictions. ``grouped=False`` retains a
-        serial THD numerical/benchmark oracle; production execution is grouped.
+        the padded Sortformer speaker predictions. Production execution is always
+        layer-synchronous and grouped; the low-level container retains a serial oracle.
         """
         if self.online_inference_enabled:
             raise RuntimeError(
                 "forward_sequence_packed is an offline API and cannot run while online_inference() is enabled."
             )
-        return self._forward_sequence_packed(
-            audio_signal, length, spk_targets, return_experts=return_experts, grouped=grouped
-        )
+        return self._forward_sequence_packed(audio_signal, length, spk_targets, return_experts=return_experts)
 
     def train(self, mode: bool = True) -> "ParallelExpertEncoder":
         """Set training mode, but keep frozen experts in eval.
@@ -1070,7 +1066,6 @@ def _forward_all_sequence_packed_training(self, audio_signal, length):
             'backend': self.sequence_packed_ggemm_backend,
             'moe_mode': self._sequence_packed_moe_execution_mode(),
             'fused_qkv': True,
-            'strict': False,
         }
         if not self.activation_checkpointing or not torch.is_grad_enabled():
             return grouped_forward(audio_signal, length, **grouped_kwargs)
@@ -1507,7 +1502,6 @@ def _forward_sequence_packed(self, audio_signal, length, spk_targets=None, *, re
                     backend=self.sequence_packed_ggemm_backend,
                     moe_mode=self._sequence_packed_moe_execution_mode(),
                     fused_qkv=True,
-                    strict=False,
                 )
 
         _validate_packed_expert_lengths(packed)

nemo/collections/asr/parts/packed_sequence.py

@@ -151,7 +151,7 @@ def _normalize_lengths(lengths: Tensor, *, batch_size: int, max_length: int, dev
         raise ValueError(f"lengths must have shape ({batch_size},), got {tuple(lengths.shape)}.")
     if lengths.device != device:
         raise ValueError(f"lengths must be on {device}, got {lengths.device}.")
-    if lengths.is_floating_point() or lengths.is_complex():
+    if lengths.dtype == torch.bool or lengths.is_floating_point() or lengths.is_complex():
         raise TypeError(f"lengths must have an integer dtype, got {lengths.dtype}.")
     lengths = lengths.to(torch.int64)
     # Varlen kernels need a host max length. Validate the same host copy so the

packed_sequence_asr_encoders_design.md

@@ -30,8 +30,8 @@ Opt-in entry points are collision-free:
 - `TransformerEncoder.forward_sequence_packed(...)` (also used by `MoETransformerEncoder`).
 - `GGEMMTransformerEncoder.forward_all_sequence_packed(...)`: serial THD oracle.
 - `GGEMMTransformerEncoder.forward_grouped_sequence_packed(...)`: layer-synchronous grouped THD.
-- `ParallelExpertEncoder.forward_sequence_packed(..., grouped=True)`: grouped production default; `grouped=False` is an
-  explicit numerical/benchmark oracle.
+- `ParallelExpertEncoder.forward_sequence_packed(...)`: layer-synchronous grouped production path; the low-level GGEMM
+  container retains the serial numerical/benchmark oracle.
 - `AudioPerceptionModule.forward_sequence_packed(...)`.
 
 Capability selection requires `supports_sequence_packed_output=True` plus the exact method. `packed_encoder_sequences`
@@ -94,8 +94,8 @@ only fields receive defaults during construction/restoration.
   the common production case.
 - Layer state stores metadata separately from initial packed data, and padded pre-encoder tensors are released before the
   Transformer loop.
-- Production uses non-strict attention bucketing, allowing old rel-pos, causal, mixed-head-dimension, and mixed-mode
-  checkpoints to use correct multiple THD buckets. `strict=True` remains an explicit diagnostic asserting one bucket.
+- Automatic attention bucketing allows old rel-pos, causal, mixed-head-dimension, and mixed-mode checkpoints to use
+  the correct number of THD buckets without exposing a diagnostic-only strict mode.
 - Unsupported custom experts keep the serial packed capability protocol. Existing legacy
   `GGEMMTransformerEncoder.forward_packed` keeps its older head-packed meaning and behavior.
 

scripts/speech_recognition/benchmark_packed_asr_encoders.py

@@ -479,14 +479,17 @@ def _run_iteration(encoder_name, model, inputs, lengths, speaker_targets, implem
 
 def _valid_output(encoder_name, model, inputs, lengths, speaker_targets, implementation):
     if encoder_name == "pee":
-        if implementation in ('serial_thd', 'grouped_thd'):
-            packed = model.forward_sequence_packed(
+        if implementation == 'serial_thd':
+            packed = model._forward_sequence_packed(
                 inputs,
                 lengths,
                 spk_targets=speaker_targets,
-                grouped=implementation == 'grouped_thd',
+                grouped=False,
             )
             return packed.data, packed.lengths
+        if implementation == 'grouped_thd':
+            packed = model.forward_sequence_packed(inputs, lengths, spk_targets=speaker_targets)
+            return packed.data, packed.lengths
         output, output_lengths = model(inputs, lengths, spk_targets=speaker_targets)
     elif implementation == 'native_thd':
         packed = model.forward_sequence_packed(inputs, lengths, bypass_pre_encode=True)

tests/collections/asr/test_packed_pee_grouped.py

@@ -65,7 +65,7 @@ def test_pee_grouped_and_serial_thd_fusion_match():
     lengths = torch.tensor([40, 17])
     targets = torch.zeros(2, 5, _N_SPK)
     with torch.no_grad():
-        serial = encoder.forward_sequence_packed(mels, lengths, spk_targets=targets, grouped=False)
+        serial = encoder._forward_sequence_packed(mels, lengths, spk_targets=targets, grouped=False)
         grouped = encoder.forward_sequence_packed(mels, lengths, spk_targets=targets)
     assert torch.equal(grouped.lengths, serial.lengths)
     assert torch.equal(grouped.cu_seqlens, serial.cu_seqlens)
@@ -240,7 +240,7 @@ def test_grouped_sequence_packed_preserves_independent_ffn_dropout_sites():
 
 
 @pytest.mark.unit
-def test_grouped_sequence_packed_non_strict_buckets_mixed_attention_features():
+def test_grouped_sequence_packed_buckets_mixed_attention_features():
     speech_config = toy_speech_expert_cfg()
     sound_config = toy_sound_expert_cfg()
     speaker_config = toy_speaker_expert_cfg()
@@ -265,7 +265,6 @@ def test_grouped_sequence_packed_non_strict_buckets_mixed_attention_features():
             signal,
             signal_lengths,
             fused_qkv=True,
-            strict=False,
         )
 
     for name in encoder.pee.expert_names:
@@ -391,18 +390,6 @@ def forward_sequence_packed(self, audio_signal, length, bypass_pre_encode=False)
         container.forward_all_sequence_packed(data, torch.tensor([3]), fused_qkv=True)
 
 
-@pytest.mark.unit
-def test_strict_grouped_sequence_packed_rejects_multiple_attention_buckets():
-    encoder = build_toy_pe_encoder().eval()
-    encoder.pee.experts['speaker'].attn_mode = 'causal'
-    signal, signal_lengths = encoder._prepare_input(
-        torch.randn(2, _MEL_FEATURES, 32),
-        torch.tensor([32, 17]),
-    )
-    with pytest.raises(ValueError, match='one compatible attention bucket'):
-        encoder.pee.forward_grouped_sequence_packed(signal, signal_lengths, strict=True)
-
-
 @pytest.mark.unit
 def test_pee_defaults_to_sparse_ragged_grouped_packed_execution():
     encoder = build_toy_pe_encoder()

tests/collections/asr/test_packed_transformer_encoder.py

@@ -84,6 +84,7 @@ def test_packed_encoder_output_all_empty_is_differentiable():
         (torch.tensor([-1, 1]), "between"),
         (torch.tensor([3, 1]), "between"),
         (torch.tensor([1.0, 1.0]), "integer dtype"),
+        (torch.tensor([True, False]), "integer dtype"),
     ],
 )
 def test_pack_encoder_output_rejects_invalid_lengths(lengths, match):