[Model] Add TP-aware MistralEncoder for FLUX.2-dev TP (#2465)

Signed-off-by: vraiti <vraiti@redhat.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

vllm_omni/diffusion/models/mistral_encoder/mistral_encoder.py

@@ -0,0 +1,667 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+TP-aware Mistral model for use as a text encoder in diffusion pipelines.
+
+Follows the same pattern as T5EncoderModel: uses vLLM's parallel linear layers
+for tensor parallelism but simple scaled_dot_product_attention instead of
+PagedAttention, so it can be used as a standalone encoder without VllmConfig.
+
+The model supports autoregressive text generation via ``generate()``
+using KV caching and the tied embedding weights as a language-model head.
+This replaces the dependency on ``Mistral3ForConditionalGeneration`` for
+caption upsampling.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PretrainedConfig
+from vllm.distributed import (
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_gather,
+)
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead, VocabParallelEmbedding
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+logger = logging.getLogger(__name__)
+
+
+class MistralRotaryEmbedding(nn.Module):
+    """RoPE implementation for the encoder. Precomputes cos/sin tables."""
+
+    def __init__(self, head_dim: int, max_position_embeddings: int, rope_theta: float = 1000000.0):
+        super().__init__()
+        inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.max_position_embeddings = max_position_embeddings
+
+    @torch.no_grad()
+    def forward(self, position_ids: torch.Tensor, dtype: torch.dtype):
+        # position_ids: (batch, seq_len)
+        inv_freq = self.inv_freq[None, :, None].float().to(position_ids.device)
+        inv_freq = inv_freq.expand(position_ids.shape[0], -1, 1)
+        pos = position_ids[:, None, :].float()
+        freqs = (inv_freq @ pos).transpose(1, 2)  # (batch, seq_len, head_dim/2)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return emb.cos().to(dtype), emb.sin().to(dtype)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(
+    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # cos/sin: (batch, seq_len, head_dim) -> (batch, 1, seq_len, head_dim)
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    if n_rep == 1:
+        return hidden_states
+    batch, num_kv_heads, slen, head_dim = hidden_states.shape
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_kv_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim)
+
+
+def _format_upsample_input(
+    prompts: list[str],
+    system_message: str,
+    images: list | None = None,
+) -> list[list[dict[str, Any]]]:
+    cleaned_txt = [p.replace("[IMG]", "") for p in prompts]
+
+    if images is None or len(images) == 0:
+        return [
+            [
+                {"role": "system", "content": [{"type": "text", "text": system_message}]},
+                {"role": "user", "content": [{"type": "text", "text": prompt}]},
+            ]
+            for prompt in cleaned_txt
+        ]
+
+    assert len(images) == len(prompts), "Number of images must match number of prompts"
+    messages = [[{"role": "system", "content": [{"type": "text", "text": system_message}]}] for _ in cleaned_txt]
+    for i, (el, batch_images) in enumerate(zip(messages, images)):
+        if batch_images is not None:
+            el.append({"role": "user", "content": [{"type": "image", "image": img} for img in batch_images]})
+        el.append({"role": "user", "content": [{"type": "text", "text": cleaned_txt[i]}]})
+    return messages
+
+
+class MistralEncoderAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.head_dim = head_dim
+
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        self.num_heads = num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        self.num_kv_heads = max(1, num_kv_heads // tp_size)
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=head_dim,
+            total_num_heads=num_heads,
+            total_num_kv_heads=num_kv_heads,
+            bias=False,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            input_size=num_heads * head_dim,
+            output_size=hidden_size,
+            bias=False,
+            prefix=f"{prefix}.o_proj",
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+        use_cache: bool = False,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
+        batch_size, seq_len, _ = hidden_states.shape
+
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        if past_key_value is not None:
+            past_k, past_v = past_key_value
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+
+        new_kv = (k, v) if use_cache else None
+
+        k = repeat_kv(k, self.num_kv_groups)
+        v = repeat_kv(v, self.num_kv_groups)
+
+        attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask, scale=self.scaling)
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(batch_size, seq_len, -1)
+        attn_output, _ = self.o_proj(attn_output)
+        return attn_output, new_kv
+
+
+class MistralEncoderMLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[intermediate_size, intermediate_size],
+            bias=False,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            bias=False,
+            prefix=f"{prefix}.down_proj",
+        )
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, _ = self.gate_up_proj(x)
+        x = self.act_fn(x)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class MistralEncoderLayer(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        intermediate_size: int,
+        rms_norm_eps: float,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.self_attn = MistralEncoderAttention(
+            hidden_size=hidden_size,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            head_dim=head_dim,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.mlp = MistralEncoderMLP(
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            prefix=f"{prefix}.mlp",
+        )
+        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+        use_cache: bool = False,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, new_kv = self.self_attn(
+            hidden_states,
+            cos,
+            sin,
+            attention_mask,
+            past_key_value,
+            use_cache,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states, new_kv
+
+
+class MistralEncoderOutput:
+    """Simple output container matching HuggingFace's interface."""
+
+    def __init__(
+        self,
+        last_hidden_state: torch.Tensor,
+        hidden_states: tuple[torch.Tensor, ...] | None = None,
+        past_key_values: list[tuple[torch.Tensor, torch.Tensor]] | None = None,
+    ):
+        self.last_hidden_state = last_hidden_state
+        self.hidden_states = hidden_states
+        self.past_key_values = past_key_values
+
+
+class MistralEncoderModel(nn.Module):
+    """
+    TP-aware Mistral encoder for use as a text encoder in diffusion pipelines.
+
+    Accepts a HuggingFace Mistral3Config (or its text_config). Uses vLLM
+    parallel layers for TP but simple SDPA for attention (no PagedAttention).
+    """
+
+    def __init__(self, config: PretrainedConfig, prefix: str = ""):
+        super().__init__()
+        self._processor = None
+        self._system_message_t2i: str | None = None
+        self._system_message_i2i: str | None = None
+        # Handle Mistral3Config (has text_config) or plain MistralConfig
+        if hasattr(config, "text_config"):
+            text_config = config.text_config
+        else:
+            text_config = config
+        self.config = text_config
+
+        self.hidden_size = text_config.hidden_size
+        self.num_heads = text_config.num_attention_heads
+        self.num_kv_heads = getattr(text_config, "num_key_value_heads", text_config.num_attention_heads)
+        self.head_dim = getattr(text_config, "head_dim", None) or (self.hidden_size // self.num_heads)
+        self.intermediate_size = text_config.intermediate_size
+        self.num_layers = text_config.num_hidden_layers
+        self.rms_norm_eps = getattr(text_config, "rms_norm_eps", 1e-5)
+        self.max_position_embeddings = getattr(text_config, "max_position_embeddings", 131072)
+        self.rope_theta = getattr(text_config, "rope_theta", 1000000.0)
+        self.vocab_size = text_config.vocab_size
+
+        tp_size = get_tensor_model_parallel_world_size()
+        logger.info(
+            "MistralEncoderModel init: hidden_size=%d, num_heads=%d, "
+            "num_kv_heads=%d, head_dim=%d, num_layers=%d, tp_size=%d",
+            self.hidden_size,
+            self.num_heads,
+            self.num_kv_heads,
+            self.head_dim,
+            self.num_layers,
+            tp_size,
+        )
+
+        # Nest modules to match HF checkpoint hierarchy:
+        #   language_model.model.embed_tokens
+        #   language_model.model.layers.X...
+        #   language_model.model.norm
+        self.language_model = nn.Module()
+        self.language_model.model = nn.Module()
+        m = self.language_model.model
+
+        m.embed_tokens = VocabParallelEmbedding(self.vocab_size, self.hidden_size)
+
+        m.layers = nn.ModuleList(
+            [
+                MistralEncoderLayer(
+                    hidden_size=self.hidden_size,
+                    num_heads=self.num_heads,
+                    num_kv_heads=self.num_kv_heads,
+                    head_dim=self.head_dim,
+                    intermediate_size=self.intermediate_size,
+                    rms_norm_eps=self.rms_norm_eps,
+                    prefix=f"language_model.model.layers.{i}",
+                )
+                for i in range(self.num_layers)
+            ]
+        )
+
+        m.norm = RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
+
+        m.rotary_emb = MistralRotaryEmbedding(self.head_dim, self.max_position_embeddings, self.rope_theta)
+
+        self.language_model.lm_head = ParallelLMHead(self.vocab_size, self.hidden_size, bias=False)
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(self.parameters()).dtype
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        output_hidden_states: bool = False,
+        use_cache: bool = False,
+        past_key_values: list[tuple[torch.Tensor, torch.Tensor]] | None = None,
+        **kwargs,
+    ) -> MistralEncoderOutput:
+        m = self.language_model.model
+        hidden_states = m.embed_tokens(input_ids)
+        seq_len = input_ids.shape[1]
+
+        # Determine position offset from cached KV length
+        past_len = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+        total_len = past_len + seq_len
+
+        # Compute position_ids from attention_mask so padded tokens get position 0
+        # and real tokens get contiguous positions starting from 0.
+        if attention_mask is not None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.clamp_(min=0)
+            position_ids = position_ids[:, -seq_len:]
+        else:
+            position_ids = torch.arange(past_len, past_len + seq_len, device=hidden_states.device).unsqueeze(0)
+
+        cos, sin = m.rotary_emb(position_ids, hidden_states.dtype)
+
+        # Build causal attention mask combined with padding mask for SDPA.
+        # Mistral is a decoder-only model, so hidden states are computed with
+        # causal (autoregressive) attention even when used as an encoder.
+        min_val = torch.finfo(hidden_states.dtype).min
+        causal_mask = torch.triu(
+            torch.full((seq_len, total_len), min_val, device=hidden_states.device, dtype=hidden_states.dtype),
+            diagonal=past_len + 1,
+        )
+        # (seq_len, total_len) -> (1, 1, seq_len, total_len)
+        sdpa_mask = causal_mask.unsqueeze(0).unsqueeze(0)
+
+        if attention_mask is not None:
+            # Combine with padding mask: (batch, 1, 1, total_len)
+            padding_mask = attention_mask[:, None, None, :].to(dtype=hidden_states.dtype)
+            padding_mask = (1.0 - padding_mask) * min_val
+            sdpa_mask = sdpa_mask + padding_mask
+
+        all_hidden_states = () if output_hidden_states else None
+        new_key_values = [] if use_cache else None
+
+        for i, layer in enumerate(m.layers):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+            past_kv = past_key_values[i] if past_key_values is not None else None
+            hidden_states, layer_kv = layer(
+                hidden_states,
+                cos,
+                sin,
+                sdpa_mask,
+                past_kv,
+                use_cache,
+            )
+            if use_cache:
+                new_key_values.append(layer_kv)
+
+        hidden_states = m.norm(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        return MistralEncoderOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            past_key_values=new_key_values,
+        )
+
+    def _compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Compute full-vocab logits from hidden states using the lm_head weight."""
+        local_logits = F.linear(
+            hidden_states,
+            self.language_model.lm_head.weight,
+        )
+        if get_tensor_model_parallel_world_size() > 1:
+            return tensor_model_parallel_all_gather(local_logits)
+        return local_logits
+
+    @staticmethod
+    def _sample(
+        logits: torch.Tensor,
+        do_sample: bool,
+        temperature: float,
+    ) -> torch.Tensor:
+        """Sample or greedily select the next token from logits. Returns (batch, 1)."""
+        if do_sample:
+            logits = logits / max(temperature, 1e-7)
+            probs = F.softmax(logits, dim=-1)
+            return torch.multinomial(probs, num_samples=1)
+        return logits.argmax(dim=-1, keepdim=True)
+
+    # TODO make common. Potentially with HF mixin
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        max_new_tokens: int = 512,
+        do_sample: bool = True,
+        temperature: float = 1.0,
+        eos_token_id: int | list[int] | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Autoregressive text generation with KV caching.
+
+        Accepts the same keyword arguments as the HuggingFace
+        ``GenerationMixin.generate`` interface used by the pipeline
+        (``pixel_values`` etc. are accepted and ignored).
+
+        Returns the full token sequence including the input prompt.
+        """
+        eos_token_id = eos_token_id or getattr(self.config, "eos_token_id", None)
+        if isinstance(eos_token_id, int):
+            eos_token_ids = {eos_token_id}
+        elif isinstance(eos_token_id, list):
+            eos_token_ids = set(eos_token_id)
+        else:
+            eos_token_ids = set()
+
+        batch_size = input_ids.shape[0]
+        device = input_ids.device
+        generated = input_ids
+
+        # Prefill ----------------------------------------------------------------
+        output = self.forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            use_cache=True,
+        )
+        past_key_values = output.past_key_values
+
+        logits = self._compute_logits(output.last_hidden_state[:, -1:, :])
+        next_token = self._sample(logits.squeeze(1), do_sample, temperature)
+        if get_tensor_model_parallel_world_size() > 1:
+            torch.distributed.broadcast(next_token, src=0)
+        generated = torch.cat([generated, next_token], dim=1)
+
+        if attention_mask is not None:
+            attention_mask = torch.cat(
+                [attention_mask, torch.ones((batch_size, 1), device=device, dtype=attention_mask.dtype)],
+                dim=1,
+            )
+
+        finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
+        if eos_token_ids:
+            eos_tensor = torch.tensor(list(eos_token_ids), device=device)
+            finished = finished | torch.isin(next_token.squeeze(-1), eos_tensor)
+
+        # Decode loop -------------------------------------------------------------
+        for _ in range(max_new_tokens - 1):
+            if finished.all():
+                break
+
+            output = self.forward(
+                input_ids=next_token,
+                attention_mask=attention_mask,
+                use_cache=True,
+                past_key_values=past_key_values,
+            )
+            past_key_values = output.past_key_values
+
+            logits = self._compute_logits(output.last_hidden_state)
+            next_token = self._sample(logits.squeeze(1), do_sample, temperature)
+            if get_tensor_model_parallel_world_size() > 1:
+                torch.distributed.broadcast(next_token, src=0)
+            generated = torch.cat([generated, next_token], dim=1)
+
+            if attention_mask is not None:
+                attention_mask = torch.cat(
+                    [attention_mask, torch.ones((batch_size, 1), device=device, dtype=attention_mask.dtype)],
+                    dim=1,
+                )
+
+            if eos_token_ids:
+                finished = finished | torch.isin(next_token.squeeze(-1), eos_tensor)
+
+        return generated
+
+    def set_processor(
+        self,
+        processor,
+        system_message_t2i: str | None = None,
+        system_message_i2i: str | None = None,
+    ) -> None:
+        self._processor = processor
+        self._system_message_t2i = system_message_t2i
+        self._system_message_i2i = system_message_i2i
+
+    @torch.no_grad()
+    def upsample_prompt(
+        self,
+        prompt: str | list[str],
+        images: list | None = None,
+        temperature: float = 0.15,
+        device: torch.device | None = None,
+        max_new_tokens: int = 512,
+        max_length: int = 2048,
+    ) -> list[str]:
+        if self._processor is None:
+            raise RuntimeError("upsample_prompt() requires a processor; call set_processor() first")
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        device = device or self.device
+
+        if images is None or len(images) == 0 or (len(images) > 0 and images[0] is None):
+            system_message = self._system_message_t2i or ""
+        else:
+            system_message = self._system_message_i2i or ""
+
+        messages_batch = _format_upsample_input(prompt, system_message, images)
+
+        inputs = self._processor.apply_chat_template(
+            messages_batch,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+            padding="max_length",
+            truncation=True,
+            max_length=max_length,
+        )
+
+        inputs["input_ids"] = inputs["input_ids"].to(device)
+        inputs["attention_mask"] = inputs["attention_mask"].to(device)
+        if "pixel_values" in inputs:
+            inputs["pixel_values"] = inputs["pixel_values"].to(device, self.dtype)
+
+        generated_ids = self.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=temperature,
+            use_cache=True,
+        )
+
+        input_length = inputs["input_ids"].shape[1]
+        generated_tokens = generated_ids[:, input_length:]
+
+        return self._processor.tokenizer.batch_decode(
+            generated_tokens,
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=True,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, weight_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        params_dict.update(self.named_buffers())
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            # Skip vision components (lm_head is needed — weights are not tied)
+            if any(name.startswith(p) for p in ("vision_tower.", "multi_modal_projector.")):
+                continue
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                if name not in params_dict:
+                    break
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                loaded_params.add(name)
+                break
+            else:
+                if name not in params_dict:
+                    logger.warning("Skipping weight %s", name)
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+                loaded_params.add(name)
+
+        total_param_bytes = sum(p.numel() * p.element_size() for p in self.parameters())
+        logger.info(
+            "MistralEncoderModel load_weights: loaded %d params, total param memory: %.2f GiB",
+            len(loaded_params),
+            total_param_bytes / (1024**3),
+        )
+        return loaded_params