Preserve SpeechLM perception checkpoint dtype

nemo/collections/speechlm2/vllm/salm/model.py

@@ -59,6 +59,9 @@
 from nemo.collections.speechlm2.vllm.salm.backends import HybridBackend, make_backend
 from nemo.collections.speechlm2.vllm.salm.config import _AUDIO_PLACEHOLDER
 
+_AUDIO_INPUT_DTYPE = torch.float32
+_PERCEPTION_DTYPE = torch.bfloat16
+
 
 @MULTIMODAL_REGISTRY.register_processor(
     NeMoSpeechLMMultiModalProcessor,
@@ -98,7 +101,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
 
         with self._mark_tower_model(vllm_config, {"audio"}):
             self.perception = _load_nemo_perception(config.perception)
-            self.perception = self.perception.to(torch.float32)
 
         self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
 
@@ -137,7 +139,7 @@ def _process_audio(self, audio_input: NeMoSpeechLMAudioInputs) -> tuple[torch.Te
         audio_signal = audio_input.audio_signal
         if isinstance(audio_signal, list):
             audio_signal = torch.stack(audio_signal, dim=0)
-        audio_signal = audio_signal.to(device=device, dtype=torch.float32)
+        audio_signal = audio_signal.to(device=device, dtype=_AUDIO_INPUT_DTYPE)
         audio_lengths = audio_input.audio_signal_length.to(device=device)
 
         with torch.no_grad():
@@ -146,8 +148,6 @@ def _process_audio(self, audio_input: NeMoSpeechLMAudioInputs) -> tuple[torch.Te
                 input_signal_length=audio_lengths,
             )
 
-        audio_embeds = audio_embeds.to(torch.bfloat16)
-
         return tuple(audio_embeds[i, : audio_embed_lens[i]] for i in range(audio_embeds.shape[0]))
 
     def embed_multimodal(self, **kwargs) -> MultiModalEmbeddings:
@@ -183,9 +183,8 @@ def get_mm_mapping(self) -> MultiModelKeys:
     # ── weight loading ──
 
     def _load_perception_weights(self, perception_weights: dict[str, torch.Tensor]) -> set[str]:
-        float32_weights = {k: v.float() for k, v in perception_weights.items()}
-        self.perception.load_state_dict(float32_weights, strict=False)
-        self.perception = self.perception.to(torch.float32)
+        self.perception = self.perception.to(_PERCEPTION_DTYPE)
+        self.perception.load_state_dict(perception_weights, strict=False)
         return {"perception." + k for k in perception_weights}
 
     @staticmethod