Address PR bot review remarks

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

Author: pzelasko

nemo/collections/asr/modules/ggemm_transformer_encoder.py

@@ -1624,7 +1624,6 @@ def _sequence_packed_grouped_attention_step(self, encs, state, layer_idx, fused_
                     fused_qkv=fused_qkv,
                 )
                 continue
-            names = [item[0] for item in group]
             hidden = torch.stack([item[1] for item in group], dim=0)
             compute_dtype = _autocast_compute_dtype(hidden)
             hidden = hidden.to(compute_dtype)
@@ -1728,7 +1727,6 @@ def _sequence_packed_grouped_attention_step(self, encs, state, layer_idx, fused_
                 layer = encs[name].layers[layer_idx]
                 state[name]['x'] = state[name]['x'] + layer.drop(output)
                 continue
-            names = [item[0] for item in group]
             hidden = torch.stack([item[1] for item in group], dim=0)
             compute_dtype = _autocast_compute_dtype(hidden)
             outputs = _grouped_linear(hidden.to(compute_dtype), [item[2] for item in group])

nemo/collections/asr/modules/parallel_expert_encoder.py

@@ -1656,9 +1656,10 @@ def forward(
         use_online = self.online_inference_enabled and self.online_inference_length > 0
         runner = self._forward_online if use_online else self._forward
         outputs, lengths, experts = runner(audio_signal=audio_signal, length=length, spk_targets=spk_targets)
+        result = (outputs, lengths)
         if return_experts:
-            return outputs, lengths, experts
-        return outputs, lengths
+            result += (experts,)
+        return result
 
     def forward_sequence_packed(
         self,

nemo/collections/speechlm2/data/salm_dataset.py

@@ -137,28 +137,29 @@ def __getitem__(self, conversations: CutSet) -> dict | None:
                 packed_audio_samples, audio_cu_seqlens, audio_lens, conversations = (
                     collate_conversation_audio_packed_fault_tolerant(conversations, self.load_audio)
                 )
+                audio_inputs = {
+                    "packed_audio_samples": packed_audio_samples,
+                    "audio_cu_seqlens": audio_cu_seqlens,
+                }
             else:
                 audios, audio_lens, conversations = collate_conversation_audio_fault_tolerant(
                     conversations, self.load_audio
                 )
+                audio_inputs = {"audios": audios}
         except Exception as e:
             logging.warning(f"Error collating conversations: {e}")
             return None
         if not conversations:
             return None
         batch = {
+            **audio_inputs,
             "audio_lens": audio_lens,
             "input_ids": left_collate_vectors([c.input_ids for c in conversations], padding_value=self.pad_id),
             "loss_mask": left_collate_vectors(
                 [getattr(c, "mask", torch.empty(0)) for c in conversations], padding_value=0
             ).to(torch.bool),
             "conversations": drop_in_memory_data(conversations),
         }
-        if self.pack_audio:
-            batch["packed_audio_samples"] = packed_audio_samples
-            batch["audio_cu_seqlens"] = audio_cu_seqlens
-        else:
-            batch["audios"] = audios
         if self.multispeaker_processor is not None:
             self.multispeaker_processor(batch)
         return batch

nemo/collections/speechlm2/modules/perception.py

@@ -222,10 +222,10 @@ def forward(
 
         # b, c, t -> b, t, c
         encoded = self.proj(encoded.transpose(1, 2))
+        result = (encoded, encoded_len)
         if return_encoder_emb:
-            return encoded, encoded_len, encoder_emb.transpose(1, 2)
-        else:
-            return encoded, encoded_len
+            result += (encoder_emb.transpose(1, 2),)
+        return result
 
     @typecheck.disable_checks()
     def forward_sequence_packed(

scripts/speech_recognition/benchmark_packed_asr_encoders.py

@@ -474,7 +474,6 @@ def _run_iteration(encoder_name, model, inputs, lengths, speaker_targets, implem
                 loss = loss + auxiliary_loss
             loss.backward()
             _clear_moe_auxiliary_loss(encoder_name, model)
-    del output
 
 
 def _valid_output(encoder_name, model, inputs, lengths, speaker_targets, implementation):

tests/collections/asr/test_packed_pee_grouped.py

@@ -19,7 +19,6 @@
 
 import nemo.collections.asr.modules.ggemm_transformer_encoder as ggemm_module
 import nemo.collections.asr.modules.transformer_encoder as transformer_module
-from nemo.collections.asr.modules.ggemm_transformer_encoder import GGEMMTransformerEncoder, _ragged_grouped_mm
 from nemo.collections.asr.parts.packed_sequence import PackedEncoderActivations, pack_encoder_output
 from tests.collections.asr.test_parallel_expert_encoder import (
     _MEL_FEATURES,
@@ -462,7 +461,7 @@ def forward_sequence_packed(self, audio_signal, length, bypass_pre_encode=False)
             return PackedEncoderActivations(audio_signal, length, cu_seqlens, int(length.max()))
 
     expert = PackedOnly()
-    container = GGEMMTransformerEncoder({'custom': expert})
+    container = ggemm_module.GGEMMTransformerEncoder({'custom': expert})
     data = torch.randn(3, 4)
     output = container.forward_all_sequence_packed(data, torch.tensor([3]))
     assert output['custom'].data is data
@@ -496,7 +495,7 @@ def test_ragged_grouped_mm_backward_handles_empty_expert_and_sum_loss(offset_val
     x = torch.randn(7, 16, device='cuda', dtype=torch.bfloat16, requires_grad=True)
     weights = [torch.randn(16, 8, device='cuda', requires_grad=True) for _ in range(3)]
     offsets = torch.tensor(offset_values, device='cuda', dtype=torch.int32)
-    output = _ragged_grouped_mm(x, offsets, weights)
+    output = ggemm_module._ragged_grouped_mm(x, offsets, weights)
     output.sum().backward()
 
     assert x.grad is not None and torch.isfinite(x.grad).all()

tests/collections/asr/test_packed_sequence_round2.py

@@ -136,6 +136,23 @@ def test_pee_packed_can_return_raw_expert_outputs_without_changing_default():
     assert experts["speaker_preds"] is not None
 
 
+def test_pee_legacy_optional_expert_return_contract_is_unchanged():
+    encoder = build_toy_pe_encoder().eval()
+    mels = torch.randn(2, _MEL_FEATURES, 24)
+    lengths = torch.tensor([24, 11])
+    targets = torch.zeros(2, 3, _N_SPK)
+
+    with torch.no_grad():
+        default = encoder(mels, lengths, spk_targets=targets)
+        with_experts = encoder(mels, lengths, spk_targets=targets, return_experts=True)
+
+    assert len(default) == 2
+    assert len(with_experts) == 3
+    torch.testing.assert_close(with_experts[0], default[0])
+    assert torch.equal(with_experts[1], default[1])
+    assert set(with_experts[2]) == {"speech", "sound", "speaker_preds"}
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="PEE sound-gradient parity requires CUDA")
 def test_pee_packed_matches_legacy_trainable_sound_gradients():
     torch.manual_seed(0)

tests/collections/speechlm2/test_perception_packed_sequence.py

@@ -61,8 +61,16 @@ def test_perception_sequence_packed_matches_legacy_and_preserves_state_dict():
 
     with torch.no_grad():
         legacy, output_lengths = perception(input_signal=features, input_signal_length=lengths)
+        legacy_with_encoder = perception(
+            input_signal=features,
+            input_signal_length=lengths,
+            return_encoder_emb=True,
+        )
         packed = perception.forward_sequence_packed(input_signal=features, input_signal_length=lengths)
 
+    assert len(legacy_with_encoder) == 3
+    torch.testing.assert_close(legacy_with_encoder[0], legacy)
+    assert torch.equal(legacy_with_encoder[1], output_lengths)
     restored = unpack_encoder_output(packed, total_length=legacy.shape[1])
     valid = torch.arange(legacy.shape[1])[None, :] < output_lengths[:, None]
     torch.testing.assert_close(restored[valid], legacy[valid], rtol=1e-5, atol=1e-6)

tests/collections/speechlm2/test_perception_packed_sequence_capabilities.py

@@ -17,7 +17,6 @@
 import pytest
 import torch
 
-from nemo.collections.asr.parts.packed_sequence import unpack_encoder_output
 from tests.collections.speechlm2.test_perception_packed_sequence import _make_perception