Added checkpoint_folder_full/modeling_nemotron_h_audio.py

checkpoint_folder_full/modeling_nemotron_h_audio.py

@@ -0,0 +1,310 @@
+"""Nemotron-Dense Audex audio-understanding model for HuggingFace inference."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from transformers.cache_utils import DynamicCache
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+from .configuration_nemotron_h_audio import NemotronDenseAudexConfig
+from .modeling_nemotron_dense import NemotronDenseForCausalLM
+
+
+class NemotronDenseAudexRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.float()
+        variance = hidden_states.pow(2).mean(dim=-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        return (self.weight.float() * hidden_states).to(input_dtype)
+
+
+class NemotronDenseAudexProjector(nn.Module):
+    """Megatron sound_projection equivalent for TP1 HF inference."""
+
+    def __init__(self, config: NemotronDenseAudexConfig):
+        super().__init__()
+        self.norm = NemotronDenseAudexRMSNorm(
+            config.audio_encoder_hidden_size,
+            eps=config.audio_projector_norm_eps,
+        )
+        self.fc1 = nn.Linear(
+            config.audio_encoder_hidden_size,
+            config.audio_projector_intermediate_size,
+            bias=False,
+        )
+        self.fc2 = nn.Linear(
+            config.audio_projector_intermediate_size,
+            config.hidden_size,
+            bias=False,
+        )
+        self.activation = config.audio_projector_activation
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.norm(hidden_states)
+        hidden_states = self.fc1(hidden_states)
+        if self.activation == "relu2":
+            hidden_states = F.relu(hidden_states).pow(2)
+        elif self.activation == "gelu":
+            hidden_states = F.gelu(hidden_states)
+        else:
+            raise ValueError(f"Unsupported audio projector activation: {self.activation}")
+        return self.fc2(hidden_states)
+
+
+def _build_qwen2_audio_encoder(audio_config: dict) -> nn.Module:
+    try:
+        from transformers.models.qwen2_audio.configuration_qwen2_audio import Qwen2AudioEncoderConfig
+        from transformers.models.qwen2_audio.modeling_qwen2_audio import Qwen2AudioEncoder
+    except Exception as exc:  # pragma: no cover - version/environment guard
+        raise ImportError(
+            "Qwen2-Audio support is required for NV-Whisper. "
+            "Install a transformers build that provides transformers.models.qwen2_audio."
+        ) from exc
+
+    cfg = Qwen2AudioEncoderConfig(**dict(audio_config))
+    return Qwen2AudioEncoder(cfg)
+
+
+class NemotronDenseAudexForConditionalGeneration(NemotronDenseForCausalLM):
+    """Nemotron-Dense CausalLM plus NV-Whisper encoder and sound projection.
+
+    State dict layout intentionally keeps the baseline LLM key names:
+    `model.*` and `lm_head.*` load exactly as in the LLM-only checkpoint.
+    New audio tensors live under `audio_encoder.*` and `audio_projector.*`.
+    """
+
+    config_class = NemotronDenseAudexConfig
+    _tp_plan = None
+    _base_model_tp_plan = None
+    base_model_tp_plan = None
+
+    def __init__(self, config: NemotronDenseAudexConfig):
+        super().__init__(config)
+        self.audio_encoder = _build_qwen2_audio_encoder(config.audio_config)
+        self.audio_projector = NemotronDenseAudexProjector(config)
+
+    def encode_audio(self, input_features: torch.Tensor) -> torch.Tensor:
+        """Encode Whisper features into LLM hidden-space audio embeddings."""
+        encoder_param = next(self.audio_encoder.parameters())
+        input_features = input_features.to(device=encoder_param.device, dtype=encoder_param.dtype)
+        encoder_outputs = self.audio_encoder(input_features=input_features, return_dict=True)
+        audio_hidden = encoder_outputs.last_hidden_state
+        projector_param = next(self.audio_projector.parameters())
+        audio_hidden = audio_hidden.to(device=projector_param.device, dtype=projector_param.dtype)
+        return self.audio_projector(audio_hidden)
+
+    def _audio_embeddings_by_sample(
+        self,
+        input_features: Optional[torch.Tensor],
+        audio_embeddings: Optional[torch.Tensor],
+        batch_size: int,
+    ) -> list[torch.Tensor]:
+        if audio_embeddings is None:
+            if input_features is None:
+                raise ValueError("input_features or audio_embeddings must be provided for audio injection")
+            if input_features.ndim == 3:
+                projected = self.encode_audio(input_features)
+                if batch_size != 1:
+                    raise ValueError(
+                        "3D input_features represent a single sample. "
+                        "Use 4D (batch, clips, mel_bins, frames) features for batched audio."
+                    )
+                return [projected.reshape(-1, projected.shape[-1])]
+            if input_features.ndim == 4:
+                bsz, clips, mel_bins, frames = input_features.shape
+                if bsz != batch_size:
+                    raise ValueError(f"input_features batch {bsz} != input_ids batch {batch_size}")
+                flat_features = input_features.reshape(bsz * clips, mel_bins, frames)
+                projected = self.encode_audio(flat_features)
+                projected = projected.reshape(bsz, clips * projected.shape[1], projected.shape[-1])
+                return [projected[idx] for idx in range(bsz)]
+            raise ValueError(f"Expected 3D or 4D input_features, got {tuple(input_features.shape)}")
+
+        if audio_embeddings.ndim == 2:
+            if batch_size != 1:
+                raise ValueError("2D audio_embeddings only support batch_size=1")
+            return [audio_embeddings.to(device=self.device)]
+        if audio_embeddings.ndim == 3:
+            if audio_embeddings.shape[0] != batch_size:
+                raise ValueError(f"audio_embeddings batch {audio_embeddings.shape[0]} != input_ids batch {batch_size}")
+            return [audio_embeddings[idx].to(device=self.device) for idx in range(batch_size)]
+        raise ValueError(f"Expected 2D or 3D audio_embeddings, got {tuple(audio_embeddings.shape)}")
+
+    def prepare_inputs_embeds(
+        self,
+        input_ids: torch.LongTensor,
+        input_features: Optional[torch.Tensor] = None,
+        audio_embeddings: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if input_ids is None:
+            raise ValueError("input_ids are required when injecting audio embeddings")
+        if self.config.sound_token_id is None:
+            raise ValueError("config.sound_token_id is required for audio embedding injection")
+
+        embed_device = self.model.embed_tokens.weight.device
+        input_ids = input_ids.to(embed_device)
+        inputs_embeds = self.model.embed_tokens(input_ids).clone()
+        audio_by_sample = self._audio_embeddings_by_sample(
+            input_features=input_features,
+            audio_embeddings=audio_embeddings,
+            batch_size=input_ids.shape[0],
+        )
+
+        for batch_idx, audio in enumerate(audio_by_sample):
+            mask = input_ids[batch_idx].to(self.device) == self.config.sound_token_id
+            mask = mask.to(embed_device)
+            expected = int(mask.sum().item())
+            if expected != audio.shape[0]:
+                raise ValueError(
+                    "Mismatch between <so_embedding> token count and projected audio tokens: "
+                    f"sample={batch_idx} placeholders={expected} audio_tokens={audio.shape[0]}"
+                )
+            inputs_embeds[batch_idx, mask] = audio.to(device=embed_device, dtype=inputs_embeds.dtype)
+        return inputs_embeds
+
+    @staticmethod
+    def _is_prefill(past_key_values) -> bool:
+        if past_key_values is None:
+            return True
+        get_len = getattr(past_key_values, "get_seq_length", None)
+        if callable(get_len):
+            return get_len() == 0
+        return len(past_key_values) == 0
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        input_features=None,
+        audio_embeddings=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        **kwargs,
+    ):
+        if self._is_prefill(past_key_values) and inputs_embeds is None and (
+            input_features is not None or audio_embeddings is not None
+        ):
+            inputs_embeds = self.prepare_inputs_embeds(
+                input_ids=input_ids,
+                input_features=input_features,
+                audio_embeddings=audio_embeddings,
+            )
+        return super().prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            position_ids=position_ids,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+    def _dense_forward_from_embeds(
+        self,
+        inputs_embeds: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values=None,
+        use_cache: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        hidden_states = inputs_embeds
+        for decoder_layer in self.model.layers:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+            )
+        hidden_states = self.model.norm(hidden_states)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, past_key_values] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        input_features: Optional[torch.Tensor] = None,
+        audio_embeddings: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values=None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ):
+        if inputs_embeds is None and self._is_prefill(past_key_values) and (
+            input_features is not None or audio_embeddings is not None
+        ):
+            inputs_embeds = self.prepare_inputs_embeds(
+                input_ids=input_ids,
+                input_features=input_features,
+                audio_embeddings=audio_embeddings,
+            )
+            input_ids = None
+        if inputs_embeds is None:
+            return super().forward(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                labels=labels,
+                use_cache=use_cache,
+                return_dict=return_dict,
+                **kwargs,
+            )
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        outputs = self._dense_forward_from_embeds(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            return_dict=return_dict,
+        )
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+        )