nvidia-nemo--speech/nemo/collections/speechlm2/models/duplex_stt_model.py

# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import copy
import os
import re

import torch
from lightning import LightningModule
from omegaconf import DictConfig
from peft import PeftModel
from torch import Tensor
from torch.distributed.fsdp import fully_shard
from torch.distributed.tensor import Replicate, Shard
from torch.distributed.tensor.parallel import (
    ColwiseParallel,
    PrepareModuleInput,
    RowwiseParallel,
    SequenceParallel,
    loss_parallel,
    parallelize_module,
)

from nemo.collections.common.tokenizers import AutoTokenizer
from nemo.collections.speechlm2.data.utils import get_pad_id
from nemo.collections.speechlm2.parts.hf_hub import HFHubMixin
from nemo.collections.speechlm2.parts.label_prep import maybe_prepend_prompt_tokens, prepare_text_and_asr_labels
from nemo.collections.speechlm2.parts.lora import maybe_install_lora
from nemo.collections.speechlm2.parts.metrics.bleu import BLEU
from nemo.collections.speechlm2.parts.metrics.empty_text import EmptyTextMetric
from nemo.collections.speechlm2.parts.metrics.results_logger import ResultsLogger
from nemo.collections.speechlm2.parts.metrics.turn_taking import TurnTakingMetrics
from nemo.collections.speechlm2.parts.metrics.wer import WER
from nemo.collections.speechlm2.parts.optim_setup import configure_optimizers, is_frozen
from nemo.collections.speechlm2.parts.pretrained import (
    load_pretrained_hf,
    maybe_load_pretrained_models,
    set_model_dict_for_partial_init,
    setup_speech_encoder,
)
from nemo.collections.speechlm2.streaming.duplex_stt_inference import DuplexSTTStreamingInference
from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
from nemo.utils import logging


def maybe_rename_llm_kwargs_for_nemotron(kwargs: dict, model_cfg) -> dict:
    """This is required because Nemotron models have a different signature than other HF models."""
    if 'Nemotron' not in model_cfg.pretrained_llm:
        return kwargs
    cache = kwargs.pop("past_key_values")
    if cache is not None:
        cache_key = model_cfg.get("cache_key", "past_key_values")
        kwargs[cache_key] = cache
    return kwargs


class DuplexSTTModel(LightningModule, HFHubMixin):
    def __init__(self, cfg: dict) -> None:
        assert isinstance(cfg, dict), (
            "You must pass the config to DuplexS2SModel as a Python dict to support hyperparameter serialization "
            f"in PTL checkpoints (we got: '{type(cfg)=}')."
        )
        super().__init__()
        self.save_hyperparameters()

        self.cfg = DictConfig(cfg)
        self.source_sample_rate = self.cfg.source_sample_rate
        self.validation_save_path = os.path.join(self.cfg.validation_save_path, "validation_logs")

        self.predict_user_text = self.cfg.get("predict_user_text", False)

        # Load LLM first
        llm = load_pretrained_hf(
            self.cfg.pretrained_llm,
            pretrained_weights=self.cfg.pretrained_weights,
            trust_remote_code=self.cfg.get("trust_remote_code", False),
        ).train()

        # Initialize tokenizer with optional special tokens from config
        tokenizer_src = self.cfg.get("tokenizer_path", None) or self.cfg.pretrained_llm
        self.tokenizer = AutoTokenizer(
            tokenizer_src,
            use_fast=True,
            bos_token=self.cfg.get("bos_token", None),
            eos_token=self.cfg.get("eos_token", None),
            pad_token=self.cfg.get("pad_token", None),
        )

        # Extract LLM components with configurable attribute names
        llm_attr_name = self.cfg.get("llm_attr_name", "model")
        self.llm = getattr(llm, llm_attr_name)
        self.lm_head = llm.lm_head

        # Extract embedding layer with configurable attribute name
        embed_tokens_attr_name = self.cfg.get("embed_tokens_attr_name", "embed_tokens")
        self.embed_tokens = getattr(self.llm, embed_tokens_attr_name)
        delattr(self.llm, embed_tokens_attr_name)

        if self.predict_user_text:
            self.asr_head = copy.deepcopy(self.lm_head)
            self.embed_asr_tokens = copy.deepcopy(self.embed_tokens)

        maybe_install_lora(self)

        # Load the pretrained streaming ASR model
        setup_speech_encoder(self, pretrained_weights=self.cfg.pretrained_weights)

        maybe_load_pretrained_models(self)

        self._use_fsdp = False
        self._use_tp = False

        # Initialize streaming inference engine
        self.streaming_inference = DuplexSTTStreamingInference(self)

    @property
    def text_vocab_size(self):
        """Return the size of the text tokenizer."""
        return self.tokenizer.vocab_size

    @property
    def text_bos_id(self) -> int:
        return self.tokenizer.bos_id

    @property
    def text_eos_id(self) -> int:
        return self.tokenizer.eos_id

    @property
    def text_pad_id(self) -> int:
        """
        Text pad ID is used as a 'blank' for frames when the model is not generating text.

        DuplexSTTModel Input/Output Format:
        - Input: User audio (speech)
        - Output: Text tokens only

        Text pad ID is used for:
        1. Frames during user speech (where the model is listening)
        2. Frames after the model completes its text response

        Example:

            flow:              |---user audio---||---assistant text---||-user audio-|
            text channel:       0000000000000000  1xxxxxxx00000000002   0000000000000
            (model output)

        Where:
        - 0 indicates PAD ID (model not generating text)
        - 1 indicates BOS ID (beginning of assistant response)
        - 2 indicates EOS ID (end of assistant response)
        - x indicates text tokens corresponding to the assistant's response

        """
        return get_pad_id(self.tokenizer)

    def forward(
        self,
        input_embeds: Tensor,
        cache=None,
    ) -> dict[str, Tensor]:
        """
        Text prediction only (audio_loss_weight=0).
        """
        kwargs = dict(inputs_embeds=input_embeds, past_key_values=cache, use_cache=cache is not None, return_dict=True)
        kwargs = maybe_rename_llm_kwargs_for_nemotron(kwargs, self.cfg)
        out = self.llm(**kwargs)

        B, T = input_embeds.shape[:2]
        text_logits = self.lm_head(out['last_hidden_state'])

        asr_logits = None
        if self.predict_user_text:
            asr_in = out['last_hidden_state']
            asr_logits = self.asr_head(asr_in)  # (B, T, asr_vocab_size)

        if not self.training:
            if self.cfg.get("inference_pad_boost", None):
                text_logits[:, :, self.text_pad_id] += self.cfg.inference_pad_boost
            if self.cfg.get("inference_bos_boost", None):
                text_logits[:, :, self.text_bos_id] += self.cfg.inference_bos_boost
            if self.cfg.get("inference_eos_boost", None):
                text_logits[:, :, self.text_eos_id] += self.cfg.inference_eos_boost

        ans = {"text_logits": text_logits}
        if self.predict_user_text:
            ans["asr_logits"] = asr_logits

        if cache is not None:
            if 'Nemotron' in self.cfg.pretrained_llm:
                cache_key = self.cfg.get("cache_key", "cache_params")
                ans["cache"] = getattr(out, cache_key, out.get(cache_key))
            else:
                ans["cache"] = out["past_key_values"]

        return ans

    def _maybe_zero_out_scale_for_asr(
        self, loss_scale: torch.Tensor, text_labels: torch.Tensor, batch: dict
    ) -> torch.Tensor:
        """
        Zero out the loss scale after text_bos_id token for ASR datasets to not penalize the agent being silent in ASR training.
        """
        if batch['task'][0] == 'asr':
            for i in range(text_labels.shape[0]):
                bos_indices = (text_labels[i] == self.text_bos_id).nonzero(as_tuple=True)
                if bos_indices[0].numel() > 0:
                    bos_idx = bos_indices[0][0].item()
                    loss_scale[i, bos_idx + 1 :, :] = 0
        return loss_scale

    def prepare_inputs(self, batch: dict):

        # Speech encoder forward pass (audio is already augmented in the dataloader)
        source_encoded, source_encoded_lens, _ = self.perception(
            input_signal=batch["source_audio"],
            input_signal_length=batch["source_audio_lens"],
            return_encoder_emb=True,
        )

        source_encoded, source_encoded_lens, target_tokens = maybe_prepend_prompt_tokens(
            batch=batch,
            embed_fn=self.embed_tokens,
            source_encoded=source_encoded,
            source_encoded_lens=source_encoded_lens,
            text_pad_id=self.text_pad_id,
        )

        if (diff := target_tokens.shape[1] - source_encoded.shape[1]) < 0:
            target_tokens = torch.cat(
                [
                    target_tokens,
                    (
                        torch.ones(source_encoded.shape[0], abs(diff), device=source_encoded.device) * self.text_pad_id
                    ).to(torch.long),
                ],
                dim=-1,
            )
        elif diff > 0:
            target_tokens = target_tokens[:, : source_encoded.shape[1]]

        inputs = prepare_text_and_asr_labels(
            batch=batch,
            target_tokens=target_tokens,
            source_encoded=source_encoded,
            cfg=self.cfg,
            text_pad_id=self.text_pad_id,
            text_bos_id=self.text_bos_id,
            text_eos_id=self.text_eos_id,
            use_tp=self._use_tp,
            device_mesh=self.device_mesh if self._use_tp else None,
        )

        source_encoded = inputs["source_encoded"]
        text_inputs = inputs["text_inputs"]
        text_labels = inputs["text_labels"]
        target_token_lens = inputs["target_token_lens"]  # Use adjusted lengths from label_prep
        asr_inputs = None
        if self.predict_user_text:
            asr_inputs = inputs["asr_inputs"]
            asr_labels = inputs["asr_labels"]

        input_embeds = self.embed_tokens(text_inputs) * self.cfg.get("duplex_text_channel_weight", 1.0)
        input_embeds.add_(source_encoded[:, :-1] * self.cfg.get("duplex_user_channel_weight", 1.0))
        if self.predict_user_text:
            asr_inputs_embeds = self.embed_asr_tokens(asr_inputs) * self.cfg.get("duplex_asr_text_weight", 1.0)
            input_embeds.add_(asr_inputs_embeds)

        seq_mask = torch.ones_like(text_labels.unsqueeze(-1), device=self.device, dtype=torch.bool)

        if self.cfg.get("mask_sequence_loss", True):
            for i in range(target_token_lens.size(0)):
                speech_end_idx = target_token_lens[i]
                seq_mask[i, speech_end_idx:, :] = 0

        loss_scale = seq_mask.clone().float()
        asr_loss_scale = seq_mask.clone().float()
        if self.cfg.get("token_loss_weight"):
            token_weights = self.cfg.token_loss_weight
            pad_weight = token_weights.get("pad", 1.0)
            bos_weight = token_weights.get("bos", 1.0)
            eos_weight = token_weights.get("eos", 1.0)
            text_weight = token_weights.get("text", 1.0)

            loss_scale = (
                torch.where(
                    text_labels.unsqueeze(-1) == self.text_pad_id,
                    pad_weight,
                    torch.where(
                        text_labels.unsqueeze(-1) == self.text_bos_id,
                        bos_weight,
                        torch.where(text_labels.unsqueeze(-1) == self.text_eos_id, eos_weight, text_weight),
                    ),
                )
                * seq_mask.float()
            )
            # Don't penalize the agent replies for ASR training
            loss_scale = self._maybe_zero_out_scale_for_asr(loss_scale, text_labels, batch)
            if self.predict_user_text:
                asr_loss_scale = (
                    torch.where(
                        asr_labels.unsqueeze(-1) == self.text_pad_id,
                        pad_weight,
                        torch.where(
                            asr_labels.unsqueeze(-1) == self.text_bos_id,
                            bos_weight,
                            torch.where(asr_labels.unsqueeze(-1) == self.text_eos_id, eos_weight, text_weight),
                        ),
                    )
                    * seq_mask.float()
                )

        ans = {
            "input_embeds": input_embeds,
            "input_lens": source_encoded_lens - 1,
            "text_labels": text_labels,
            "loss_scale": loss_scale,
            "seq_mask": seq_mask,
        }
        if self.predict_user_text:
            ans["asr_labels"] = asr_labels
            ans["asr_loss_scale"] = asr_loss_scale
        return ans

    def training_step(self, batch: dict, batch_idx: int):
        for m in (self.perception.preprocessor, self.perception.encoder, self.llm):
            if is_frozen(m):
                m.eval()

        res = {
            "learning_rate": torch.as_tensor(
                self.trainer.optimizers[0].param_groups[0]['lr'] if self._trainer is not None else 0
            )
        }

        if batch["audio_data"] is not None:
            inputs = self.prepare_inputs(batch["audio_data"])

            forward_outputs = self(inputs["input_embeds"])

            num_frames = inputs["input_lens"].sum()

            with loss_parallel():
                text_logits = forward_outputs["text_logits"]
                asr_logits = None
                if self.predict_user_text:
                    asr_logits = forward_outputs["asr_logits"]

                if self.cfg.get("mask_sequence_loss", True):
                    text_logits = text_logits * inputs["seq_mask"][:, :, 0].unsqueeze(-1)

                text_loss = (
                    torch.nn.functional.cross_entropy(
                        text_logits.flatten(0, 1),
                        inputs["text_labels"].flatten(0, 1),
                        reduction="none",
                    )
                    * inputs["loss_scale"][:, :, 0].flatten(0, 1)
                ).sum(-1) / num_frames

                asr_loss = None
                if self.predict_user_text:
                    asr_loss = (
                        torch.nn.functional.cross_entropy(
                            asr_logits.flatten(0, 1),
                            inputs["asr_labels"].flatten(0, 1),
                            reduction="none",
                        )
                        * inputs["asr_loss_scale"][:, :, 0].flatten(0, 1)
                    ).sum(-1) / num_frames

                with torch.no_grad():
                    predicted_tokens = torch.argmax(text_logits, dim=-1)  # (B, T)
                    target_tokens = inputs["text_labels"]  # (B, T)
                    valid_mask = target_tokens != self.text_pad_id

                    correct_predictions = (predicted_tokens == target_tokens) & valid_mask

                    if valid_mask.sum() > 0:
                        token_accuracy = correct_predictions.sum().float() / valid_mask.sum().float()
                    else:
                        token_accuracy = torch.tensor(0.0, device=text_logits.device)

                loss = self.cfg.text_loss_weight * text_loss

                if self.predict_user_text:
                    loss = loss + self.cfg.get('asr_loss_weight', 1.0) * asr_loss

                B, T = inputs["input_embeds"].shape[:2]
                ans = {
                    "audio_loss": loss,
                    "audio_to_text_loss": text_loss,
                    "batch": B,
                    "length": T,
                    "token_accuracy": token_accuracy,
                }
                if self.predict_user_text:
                    ans["asr_loss"] = asr_loss

                res.update(ans)

        if batch["text_data"] is not None:
            text_input_ids = batch["text_data"]["text_tokens"][:, :-1]
            text_target = batch["text_data"]["text_tokens"][:, 1:]

            text_out = self.llm(
                inputs_embeds=self.embed_tokens(text_input_ids),
                past_key_values=None,
                use_cache=False,
                return_dict=True,
            )
            text_logits = self.lm_head(text_out['last_hidden_state'])  # (B, T, Vt)

            text_loss = torch.nn.functional.cross_entropy(
                text_logits.flatten(0, 1),  # (B, T, Vt) -> (*, Vt)
                text_target.flatten(0, 1),
                ignore_index=self.text_pad_id,
            )
            res.update(
                {
                    "text_to_text_loss": text_loss,
                }
            )

        res["loss"] = (1.0 - self.cfg.get('text_to_text_loss_weight', 0.0)) * res.get(
            "audio_loss", 0.0
        ) + self.cfg.get('text_to_text_loss_weight', 0.0) * res.get("text_to_text_loss", 0.0)
        self.log_dict(res, on_step=True)

        return res

    def on_validation_epoch_start(self) -> None:
        self.results_logger = ResultsLogger(self.validation_save_path).reset()
        self.bleu = BLEU().reset()

        self.turn_taking_metrics = TurnTakingMetrics(
            eos_token_id=self.text_eos_id,
            bos_token_id=self.text_bos_id,
            tolerance=13,
            latency_multiplier=0.08,
        ).reset()

        if self.predict_user_text:
            self.src_bleu = BLEU().reset()
            self.src_wer = WER().reset()
            self.empty_user_text = EmptyTextMetric().reset()

    def on_validation_epoch_end(self, prefix="val") -> None:
        bleu = self.bleu.compute()
        for k, m in bleu.items():
            if "qa" not in k and "mmsu" not in k:
                self.log(f"{prefix}_{k}", m.to(self.device), on_epoch=True, sync_dist=True)

        acc_metrics = self.results_logger.compute_and_save()

        for name, result_dict in acc_metrics.items():
            if 'acc' in result_dict:
                self.log(f"{prefix}_{name}_acc", result_dict['acc'].to(self.device), on_epoch=True, sync_dist=True)

            if 'mcq_acc' in result_dict:
                self.log(
                    f"{prefix}_{name}_mcq_acc", result_dict['mcq_acc'].to(self.device), on_epoch=True, sync_dist=True
                )

        turn_taking_metrics = self.turn_taking_metrics.compute()
        for k, m in turn_taking_metrics.items():
            self.log(f"{prefix}_{k}", m.to(self.device), on_epoch=True, sync_dist=True)

        if self.predict_user_text:
            src_bleu = self.src_bleu.compute()
            for k, m in src_bleu.items():
                self.log(f"{prefix}_src_{k}", m.to(self.device), on_epoch=True, sync_dist=True)
            src_wer = self.src_wer.compute()
            for k, m in src_wer.items():
                self.log(f"{prefix}_src_{k}", m.to(self.device), on_epoch=True, sync_dist=True)
            empty_user_text = self.empty_user_text.compute()
            for k, m in empty_user_text.items():
                self.log(f"{prefix}_src_{k}", m.to(self.device), on_epoch=True, sync_dist=True)

        torch.cuda.empty_cache()
        torch.cuda.synchronize()

    def validation_step(self, batch: dict, batch_idx: int):
        for name, dataset_batch in batch.items():
            if dataset_batch is None:
                continue

            dataset_batch = dataset_batch["audio_data"]

            prompt_tokens = dataset_batch.get("prompt_tokens", None)
            prompt_token_lens = dataset_batch.get("prompt_token_lens", None)

            results = self.streaming_inference.offline_inference(
                dataset_batch["source_audio"],
                dataset_batch["source_audio_lens"],
                prompt_tokens=prompt_tokens,
                prompt_token_lens=prompt_token_lens,
            )

            # Strip timestamps for metrics
            text_clean = [re.sub(r"<[\|$].*?[\|$]>", "", s).strip() for s in results["text"]]

            # Agent text metrics
            self.bleu.update(name=name, refs=dataset_batch["target_texts"], hyps=text_clean)
            if "source_tokens" in dataset_batch and results["tokens_text"] is not None:
                self.turn_taking_metrics.update(
                    name=name, source_tokens=dataset_batch["source_tokens"], pred_tokens=results["tokens_text"]
                )

            # User text metrics
            if self.predict_user_text:
                self.src_bleu.update(name=name, refs=dataset_batch["source_texts"], hyps=results["src_text"])
                self.src_wer.update(name=name, refs=dataset_batch["source_texts"], hyps=results["src_text"])
                self.empty_user_text.update(name=name, hyps=results["src_text"])

            self.results_logger.update(
                name=name,
                refs=dataset_batch["target_texts"],
                hyps=results["text"],
                samples_id=dataset_batch['sample_id'],
                user_audio=dataset_batch["source_audio"],
                user_audio_sr=self.source_sample_rate,
                src_refs=dataset_batch["source_texts"],
                src_hyps=results["src_text"],
            )

    def on_test_epoch_start(self) -> None:
        return self.on_validation_epoch_start()

    def on_test_epoch_end(self) -> None:
        return self.on_validation_epoch_end(prefix="test")

    def test_step(self, *args, **kwargs):
        return self.validation_step(*args, **kwargs)

    def predict_step(self, batch: dict, batch_idx: int, dataloader_idx: int = 0):
        batch = batch["audio_data"]

        prompt_tokens = batch.get("prompt_tokens", None)
        prompt_token_lens = batch.get("prompt_token_lens", None)

        prediction = self.streaming_inference.offline_inference(
            batch["source_audio"],
            batch["source_audio_lens"],
            input_pad_len=self.cfg.prediction.max_new_seconds * self.cfg.prediction.input_sample_rate,
            prompt_tokens=prompt_tokens,
            prompt_token_lens=prompt_token_lens,
        )
        prediction["sample_id"] = batch["sample_id"]
        return prediction

    def _get_bos_embedding(self) -> torch.Tensor:
        """Get BOS embedding for AR decoding."""
        text_bos = torch.full((1,), fill_value=self.text_pad_id, device=self.device)
        input_embeds = self.embed_tokens(text_bos)
        return input_embeds

    def _get_asr_bos_embedding(self) -> torch.Tensor:
        """Get ASR BOS embedding for AR decoding."""
        text_bos = torch.full((1,), fill_value=self.text_pad_id, device=self.device)
        input_embeds = self.embed_asr_tokens(text_bos)
        return input_embeds

    def backward(self, *args, **kwargs):
        with loss_parallel():
            super().backward(*args, **kwargs)

    def configure_optimizers(self):
        return configure_optimizers(self)

    @property
    def oomptimizer_schema(self) -> dict:
        """
        Return a typing schema for optimal batch size calibration.
        """
        return {
            "cls": dict,
            "inputs": [
                {"name": "source_audio", "type": NeuralType(("B", "T"), AudioSignal()), "seq_length": "input"},
                {"name": "source_audio_lens", "type": NeuralType(("B",), LengthsType()), "seq_length": "input"},
                {
                    "name": "target_tokens",
                    "type": NeuralType(("B", "T"), LabelsType()),
                    "seq_length": "output",
                    "vocab_size": self.tokenizer.vocab_size,
                },
            ],
        }

    def configure_model(self) -> None:
        device_mesh = self.device_mesh
        if device_mesh is None:
            return

        llm = self.llm
        if isinstance(llm, PeftModel):
            llm = llm.base_model.model

        if (tp_mesh := device_mesh["tensor_parallel"]).size() > 1:
            self._use_tp = True

            plan = {
                "layers.0": PrepareModuleInput(
                    input_layouts=(Replicate(),),
                    desired_input_layouts=(Shard(1),),
                    use_local_output=True,
                ),
                "norm": SequenceParallel(),
            }
            parallelize_module(llm, tp_mesh, plan)

            for transformer_block in llm.layers:
                plan = {
                    "input_layernorm": SequenceParallel(),
                    "self_attn.q_proj": ColwiseParallel(),
                    "self_attn.k_proj": ColwiseParallel(),
                    "self_attn.v_proj": ColwiseParallel(),
                    "self_attn.o_proj": RowwiseParallel(output_layouts=Shard(1)),
                    "post_attention_layernorm": SequenceParallel(),
                    "mlp": PrepareModuleInput(
                        input_layouts=(Shard(1),),
                        desired_input_layouts=(Replicate(),),
                    ),
                    "mlp.gate_proj": ColwiseParallel(),
                    "mlp.up_proj": ColwiseParallel(),
                    "mlp.down_proj": RowwiseParallel(output_layouts=Shard(1)),
                }

                attn_layer = transformer_block.self_attn

                try:
                    config = self.llm.config

                    num_attention_heads = getattr(config, 'num_attention_heads', None)
                    num_key_value_heads = getattr(config, 'num_key_value_heads', None)
                    hidden_size = getattr(config, 'hidden_size', None)

                    if all([num_attention_heads, num_key_value_heads, hidden_size]):
                        for attr_name, val in [
                            ("num_attention_heads", num_attention_heads),
                            ("num_key_value_heads", num_key_value_heads),
                            ("hidden_size", hidden_size),
                        ]:
                            if val % tp_mesh.size() != 0:
                                logging.warning(
                                    f"config.{attr_name}={val} is not divisible by {tp_mesh.size()=}: "
                                    f"set a different tensor parallelism size to avoid errors."
                                )

                        if hasattr(attn_layer, 'num_heads'):
                            attn_layer.num_heads = num_attention_heads // tp_mesh.size()
                        elif hasattr(attn_layer, 'num_attention_heads'):
                            attn_layer.num_attention_heads = num_attention_heads // tp_mesh.size()

                        if hasattr(attn_layer, 'num_key_value_heads'):
                            attn_layer.num_key_value_heads = num_key_value_heads // tp_mesh.size()

                        if hasattr(attn_layer, 'hidden_size'):
                            attn_layer.hidden_size = hidden_size // tp_mesh.size()

                        logging.info(
                            f"Configured tensor parallel for attention: "
                            f"heads={num_attention_heads // tp_mesh.size()}, "
                            f"kv_heads={num_key_value_heads // tp_mesh.size()}, "
                            f"hidden_size={hidden_size // tp_mesh.size()}"
                        )
                    else:
                        raise AttributeError("Required config attributes not found")

                except Exception as e:
                    logging.warning(f"Failed to configure tensor parallel using config: {e}")
                    logging.warning("Falling back to attention layer attributes...")

                    try:
                        for attr in ("num_heads", "num_key_value_heads", "hidden_size"):
                            if hasattr(attn_layer, attr):
                                val = getattr(attn_layer, attr)
                                if val % tp_mesh.size() != 0:
                                    logging.warning(
                                        f"attn_layer.{attr}={val} is not divisible by {tp_mesh.size()=}: "
                                        f"set a different tensor parallelism size to avoid errors."
                                    )
                                setattr(attn_layer, attr, val // tp_mesh.size())
                    except Exception as fallback_e:
                        logging.warning(f"Both config and fallback methods failed: {fallback_e}")
                        logging.warning("Skipping tensor parallel configuration for this attention layer")

            for m in (self.lm_head,):
                parallelize_module(
                    m,
                    tp_mesh,
                    ColwiseParallel(
                        input_layouts=Shard(1),
                        output_layouts=Shard(-1),
                        use_local_output=False,
                    ),
                )

        if (dp_mesh := device_mesh["data_parallel"]).size() > 1:
            assert dp_mesh.ndim == 1
            self._use_fsdp = True

            fsdp_config = {"mesh": dp_mesh}

            for idx, layer in enumerate(llm.layers):
                llm.layers[idx] = fully_shard(layer, **fsdp_config)
            self.embed_tokens = fully_shard(self.embed_tokens, **fsdp_config)
            self.llm = fully_shard(self.llm, **fsdp_config)
            self.lm_head = fully_shard(self.lm_head, **fsdp_config)
            self.perception = fully_shard(self.perception, **fsdp_config)
            if self.predict_user_text:
                self.asr_head = fully_shard(self.asr_head, **fsdp_config)
                self.embed_asr_tokens = fully_shard(self.embed_asr_tokens, **fsdp_config)

    def load_state_dict(self, state_dict, strict: bool = True):
        try:
            return super().load_state_dict(state_dict, strict=strict)
        except RuntimeError:
            logging.info("Error loading model state_dict !! Retrying with partial initialization!")
            model_dict = set_model_dict_for_partial_init(state_dict, self.state_dict())
            return super().load_state_dict(model_dict, strict=False)