refactor(speechlm2): scope portable PEE support to serving

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

Author: pzelasko

nemo/collections/asr/models/sortformer_diar_models.py

@@ -1004,7 +1004,7 @@ def forward_streaming_step(
             spkcache_fifo_chunk_pre_encode_embs = self.sortformer_modules.concat_embs(
                 [streaming_state.spkcache, streaming_state.fifo, chunk_pre_encode_embs],
                 dim=1,
-                # A Sortformer nested in phPEE can retain Lightning's cached
+                # A Sortformer nested in ParallelExpertEncoder can retain Lightning's cached
                 # construction device after its parent moves. Activations are
                 # authoritative for streaming lengths and masks.
                 device=chunk_pre_encode_embs.device,

nemo/collections/asr/modules/parallel_expert_encoder.py

@@ -278,6 +278,14 @@ def __init__(self, cfg: DictConfig, trainer: Optional[Trainer] = None):
             chunk_size_seconds=self._cfg.get("chunk_size_seconds", None),
             sync_max_audio_length=self._cfg.get("sync_max_audio_length", False),
         )
+        # Preserve the architecture-only bundle configuration for consolidated
+        # SpeechLM checkpoint export and serving reconstruction.
+        self.encoder._bundle_config = _clone_config(self._cfg)
+        self.encoder._bundle_config.diar_normalize_type = self.encoder.diar_normalize_type
+        self.encoder._bundle_config.speaker_feature_config_version = _SPEAKER_FEATURE_CONFIG_VERSION
+        self.encoder._bundle_config.speaker_feature_mode = self.encoder.speaker_feature_mode
+        self.encoder._bundle_config.speaker_activity_threshold = self.encoder.speaker_activity_threshold
+        self.encoder._bundle_config.sync_max_audio_length = self.encoder.sync_max_audio_length
 
     @staticmethod
     def _validate_bundle_schema(cfg: DictConfig) -> None:
@@ -375,6 +383,21 @@ def load_from_nemo(
         )
         return bundle.encoder
 
+    @classmethod
+    def from_inline_config(
+        cls,
+        cfg: Union[DictConfig, dict],
+        *,
+        map_location: Union[str, torch.device] = "cpu",
+    ) -> ParallelExpertEncoder:
+        """Construct the encoder architecture without loading standalone weights.
+
+        Consolidated SpeechLM checkpoints supply the encoder tensors from their
+        root state dictionary after constructing it from this embedded config.
+        """
+        shell = cls(cfg=OmegaConf.create(cfg), trainer=None)
+        return shell.encoder.to(map_location)
+
     @classmethod
     def save_to_nemo(
         cls,
@@ -445,6 +468,7 @@ class ParallelExpertEncoder(nn.Module):
     :class:`TransformerEncoder` used by Transformer AED ASR checkpoints.
     """
 
+    supports_external_speaker_targets = True
     supports_sequence_packed_output = True
 
     def __init__(

tests/collections/speechlm2/test_vllm_plugin.py

@@ -55,8 +55,6 @@ def test_full_stack_asr_exports_grouped_encoder_dependencies():
 from nemo.collections.asr.models import ASRModel, SortformerEncLabelModel
 from nemo.collections.asr.modules import (
     ConvASRDecoder,
-    GGEMMTransformerEncoder,
-    MoETransformerEncoder,
     ParallelExpertEncoder,
     TransformerEncoder,
 )
@@ -81,8 +79,6 @@ def test_full_stack_asr_exports_grouped_encoder_dependencies():
     for dependency in (
         SortformerEncLabelModel,
         ConvASRDecoder,
-        GGEMMTransformerEncoder,
-        MoETransformerEncoder,
         TransformerEncoder,
         ParallelExpertEncoder,
         ParallelExpertEncoderPT,