qwen3_decoder.py

# coding=utf-8
from __future__ import annotations

import json
from pathlib import Path
from typing import Optional

import torch
import torch.nn as nn
import torch.utils.checkpoint
from safetensors.torch import load_file
from transformers.activations import ACT2FN
from transformers.modeling_outputs import BaseModelOutputWithPast

from .gpt2_decoder import PackedSequenceMetadata, MossTTSNanoGPT2Model

try:
    from flash_attn import flash_attn_func, flash_attn_varlen_func
    from flash_attn.bert_padding import pad_input, unpad_input

    _FLASH_ATTN_AVAILABLE = True
except Exception:
    flash_attn_func = None
    flash_attn_varlen_func = None
    pad_input = None
    unpad_input = None
    _FLASH_ATTN_AVAILABLE = False


class MossQwen3RMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)


class MossQwen3RotaryEmbedding(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()
        head_dim = int(getattr(config, "head_dim", config.hidden_size // config.num_attention_heads))
        rope_theta = getattr(config, "rope_theta", None)
        if rope_theta is None:
            rope_scaling = getattr(config, "rope_scaling", None)
            if isinstance(rope_scaling, dict):
                rope_theta = rope_scaling.get("rope_theta")
        if rope_theta is None:
            rope_theta = 1000000.0
        rope_theta = float(rope_theta)
        self.head_dim = head_dim
        self.rope_theta = rope_theta
        self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False)

    def _compute_inv_freq(self, device: Optional[torch.device] = None) -> torch.Tensor:
        return 1.0 / (
            self.rope_theta ** (torch.arange(0, self.head_dim, 2, device=device, dtype=torch.float32) / self.head_dim)
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_ids: torch.LongTensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        inv_freq = self._compute_inv_freq(device=hidden_states.device)
        freqs = torch.einsum(
            "bs,d->bsd",
            position_ids.to(device=hidden_states.device, dtype=inv_freq.dtype),
            inv_freq,
        )
        emb = torch.cat((freqs, freqs), dim=-1)
        return emb.cos().to(dtype=hidden_states.dtype), emb.sin().to(dtype=hidden_states.dtype)


def rotate_half(hidden_states: torch.Tensor) -> torch.Tensor:
    first_half = hidden_states[..., : hidden_states.shape[-1] // 2]
    second_half = hidden_states[..., hidden_states.shape[-1] // 2 :]
    return torch.cat((-second_half, first_half), dim=-1)


def apply_rotary_pos_emb(
    query: torch.Tensor,
    key: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    cos = cos.unsqueeze(-2)
    sin = sin.unsqueeze(-2)
    query = (query * cos) + (rotate_half(query) * sin)
    key = (key * cos) + (rotate_half(key) * sin)
    return query, key


def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    if n_rep == 1:
        return hidden_states
    batch, seq_len, num_key_value_heads, head_dim = hidden_states.shape
    hidden_states = hidden_states[:, :, :, None, :].expand(batch, seq_len, num_key_value_heads, n_rep, head_dim)
    return hidden_states.reshape(batch, seq_len, num_key_value_heads * n_rep, head_dim)


class MossQwen3MLP(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.act_fn = ACT2FN[getattr(config, "hidden_act", "silu")]

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))


class MossQwen3Attention(nn.Module):
    def __init__(self, config, layer_idx: int) -> None:
        super().__init__()
        self.config = config
        self.layer_idx = int(layer_idx)
        self.hidden_size = int(config.hidden_size)
        self.num_heads = int(config.num_attention_heads)
        self.num_key_value_heads = int(config.num_key_value_heads)
        self.head_dim = int(getattr(config, "head_dim", self.hidden_size // self.num_heads))
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.scaling = self.head_dim ** -0.5
        self.attention_dropout = float(getattr(config, "attention_dropout", 0.0))
        self.attn_implementation = str(getattr(config, "_attn_implementation", "eager"))

        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=bool(config.attention_bias))
        self.k_proj = nn.Linear(
            self.hidden_size,
            self.num_key_value_heads * self.head_dim,
            bias=bool(config.attention_bias),
        )
        self.v_proj = nn.Linear(
            self.hidden_size,
            self.num_key_value_heads * self.head_dim,
            bias=bool(config.attention_bias),
        )
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=bool(config.attention_bias))
        self.q_norm = MossQwen3RMSNorm(self.head_dim, eps=float(config.rms_norm_eps))
        self.k_norm = MossQwen3RMSNorm(self.head_dim, eps=float(config.rms_norm_eps))

    def _causal_attention_mask(
        self,
        attention_mask: Optional[torch.Tensor],
        query_length: int,
        key_length: int,
        device: torch.device,
    ) -> torch.Tensor:
        query_positions = torch.arange(query_length, device=device, dtype=torch.long)
        query_positions = query_positions + max(key_length - query_length, 0)
        key_positions = torch.arange(key_length, device=device, dtype=torch.long)
        causal = key_positions.unsqueeze(0) <= query_positions.unsqueeze(1)
        causal = causal.unsqueeze(0).unsqueeze(0)
        if attention_mask is None:
            return causal
        key_mask = attention_mask[:, None, None, :].to(dtype=torch.bool)
        return causal & key_mask

    def _eager_attention(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attention_mask: Optional[torch.Tensor],
    ) -> torch.Tensor:
        key = repeat_kv(key, self.num_key_value_groups)
        value = repeat_kv(value, self.num_key_value_groups)
        query = query.transpose(1, 2)
        key = key.transpose(1, 2)
        value = value.transpose(1, 2)
        scores = torch.matmul(query, key.transpose(-1, -2)) * self.scaling
        mask = self._causal_attention_mask(
            attention_mask=attention_mask,
            query_length=query.shape[-2],
            key_length=key.shape[-2],
            device=query.device,
        )
        scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
        probs = torch.softmax(scores, dim=-1)
        if self.training and self.attention_dropout > 0:
            probs = torch.dropout(probs, self.attention_dropout, train=True)
        output = torch.matmul(probs, value)
        return output.transpose(1, 2).contiguous()

    def _sdpa_attention(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attention_mask: Optional[torch.Tensor],
    ) -> torch.Tensor:
        key = repeat_kv(key, self.num_key_value_groups)
        value = repeat_kv(value, self.num_key_value_groups)
        query = query.transpose(1, 2)
        key = key.transpose(1, 2)
        value = value.transpose(1, 2)
        mask = None
        if attention_mask is not None or query.shape[-2] != key.shape[-2]:
            mask = self._causal_attention_mask(
                attention_mask=attention_mask,
                query_length=query.shape[-2],
                key_length=key.shape[-2],
                device=query.device,
            )
        output = torch.nn.functional.scaled_dot_product_attention(
            query,
            key,
            value,
            attn_mask=mask,
            dropout_p=self.attention_dropout if self.training else 0.0,
            is_causal=mask is None,
            scale=self.scaling,
        )
        return output.transpose(1, 2).contiguous()

    def _flash_attention(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attention_mask: Optional[torch.Tensor],
        packed_metadata: Optional[PackedSequenceMetadata],
    ) -> torch.Tensor:
        if not _FLASH_ATTN_AVAILABLE:
            raise ImportError("flash_attn is not installed, but attn_implementation='flash_attention_2' was requested.")
        if query.device.type != "cuda":
            raise ValueError("flash_attention_2 requires CUDA tensors.")
        if query.dtype not in (torch.float16, torch.bfloat16):
            raise ValueError(f"flash_attention_2 requires fp16/bf16 tensors, got dtype={query.dtype}.")

        dropout_p = self.attention_dropout if self.training else 0.0
        if packed_metadata is not None:
            if packed_metadata.indices is not None:
                query = query.reshape(-1, self.num_heads, self.head_dim).index_select(0, packed_metadata.indices)
                key = key.reshape(-1, self.num_key_value_heads, self.head_dim).index_select(0, packed_metadata.indices)
                value = value.reshape(-1, self.num_key_value_heads, self.head_dim).index_select(0, packed_metadata.indices)
            output = flash_attn_varlen_func(
                query,
                key,
                value,
                packed_metadata.cu_seqlens,
                packed_metadata.cu_seqlens,
                packed_metadata.max_seqlen,
                packed_metadata.max_seqlen,
                dropout_p=dropout_p,
                causal=True,
            )
            if packed_metadata.indices is None:
                return output
            return pad_input(
                output,
                packed_metadata.indices,
                packed_metadata.batch_size,
                packed_metadata.seq_len,
            )

        if attention_mask is None or bool(attention_mask.all()):
            return flash_attn_func(query, key, value, dropout_p=dropout_p, causal=True)

        if query.shape[1] != key.shape[1]:
            query_attention_mask = attention_mask[:, -query.shape[1] :]
            unpadded_query, query_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(
                query,
                query_attention_mask,
            )
            unpadded_key, _, cu_seqlens_k, max_seqlen_k, _ = unpad_input(key, attention_mask)
            unpadded_value, _, _, _, _ = unpad_input(value, attention_mask)
            output = flash_attn_varlen_func(
                unpadded_query,
                unpadded_key,
                unpadded_value,
                cu_seqlens_q,
                cu_seqlens_k,
                max_seqlen_q,
                max_seqlen_k,
                dropout_p=dropout_p,
                causal=True,
            )
            return pad_input(output, query_indices, query.shape[0], query.shape[1])

        unpadded_query, indices, cu_seqlens, max_seqlen, _ = unpad_input(query, attention_mask)
        unpadded_key, _, _, _, _ = unpad_input(key, attention_mask)
        unpadded_value, _, _, _, _ = unpad_input(value, attention_mask)
        output = flash_attn_varlen_func(
            unpadded_query,
            unpadded_key,
            unpadded_value,
            cu_seqlens,
            cu_seqlens,
            max_seqlen,
            max_seqlen,
            dropout_p=dropout_p,
            causal=True,
        )
        return pad_input(output, indices, query.shape[0], query.shape[1])

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor] = None,
        packed_metadata: Optional[PackedSequenceMetadata] = None,
        layer_past: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
    ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:
        input_shape = hidden_states.shape[:-1]
        query_states = self.q_norm(
            self.q_proj(hidden_states).view(*input_shape, self.num_heads, self.head_dim)
        )
        key_states = self.k_norm(
            self.k_proj(hidden_states).view(*input_shape, self.num_key_value_heads, self.head_dim)
        )
        value_states = self.v_proj(hidden_states).view(*input_shape, self.num_key_value_heads, self.head_dim)

        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        if layer_past is not None:
            past_key, past_value = layer_past
            key_states = torch.cat([past_key.to(device=key_states.device, dtype=key_states.dtype), key_states], dim=1)
            value_states = torch.cat(
                [past_value.to(device=value_states.device, dtype=value_states.dtype), value_states],
                dim=1,
            )

        present = (key_states, value_states) if use_cache else None
        if self.attn_implementation == "flash_attention_2":
            attn_output = self._flash_attention(
                query=query_states,
                key=key_states,
                value=value_states,
                attention_mask=attention_mask,
                packed_metadata=packed_metadata,
            )
        elif self.attn_implementation == "sdpa":
            attn_output = self._sdpa_attention(
                query=query_states,
                key=key_states,
                value=value_states,
                attention_mask=attention_mask,
            )
        else:
            attn_output = self._eager_attention(
                query=query_states,
                key=key_states,
                value=value_states,
                attention_mask=attention_mask,
            )
        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        return self.o_proj(attn_output), present


class MossQwen3DecoderLayer(nn.Module):
    def __init__(self, config, layer_idx: int) -> None:
        super().__init__()
        self.self_attn = MossQwen3Attention(config=config, layer_idx=layer_idx)
        self.mlp = MossQwen3MLP(config)
        self.input_layernorm = MossQwen3RMSNorm(config.hidden_size, eps=float(config.rms_norm_eps))
        self.post_attention_layernorm = MossQwen3RMSNorm(config.hidden_size, eps=float(config.rms_norm_eps))

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        packed_metadata: Optional[PackedSequenceMetadata] = None,
        layer_past: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
    ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        attn_output, present = self.self_attn(
            hidden_states=hidden_states,
            position_embeddings=position_embeddings,
            attention_mask=attention_mask,
            packed_metadata=packed_metadata,
            layer_past=layer_past,
            use_cache=use_cache,
        )
        hidden_states = residual + attn_output

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = residual + self.mlp(hidden_states)
        return hidden_states, present


class MossQwen3Model(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()
        self.config = config
        self.attn_implementation = str(getattr(config, "_attn_implementation", "eager"))
        self.padding_idx = getattr(config, "pad_token_id", None)
        self.vocab_size = int(config.vocab_size)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList(
            [MossQwen3DecoderLayer(config, layer_idx=index) for index in range(config.num_hidden_layers)]
        )
        self.norm = MossQwen3RMSNorm(config.hidden_size, eps=float(config.rms_norm_eps))
        self.rotary_emb = MossQwen3RotaryEmbedding(config)
        self.gradient_checkpointing = False
        self.gradient_checkpointing_use_reentrant = bool(
            getattr(config, "gradient_checkpointing_use_reentrant", False)
        )
        self._reset_parameters()

    def _reset_parameters(self) -> None:
        init_std = float(getattr(self.config, "initializer_range", 0.02))
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=init_std)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, mean=0.0, std=init_std)

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def load_qwen3_pretrained_weights(self, pretrained_path: str) -> None:
        model_dir = Path(pretrained_path)
        index_path = model_dir / "model.safetensors.index.json"
        if not index_path.exists():
            raise FileNotFoundError(f"Missing Qwen3 safetensors index: {index_path}")
        with index_path.open("r", encoding="utf-8") as handle:
            index = json.load(handle)
        weight_map = index.get("weight_map", {})
        shard_to_keys: dict[str, list[str]] = {}
        for key, shard in weight_map.items():
            if not key.startswith("model."):
                continue
            shard_to_keys.setdefault(str(shard), []).append(key)

        state_dict = self.state_dict()
        loaded_state = {}
        for shard, keys in sorted(shard_to_keys.items()):
            shard_tensors = load_file(str(model_dir / shard), device="cpu")
            for key in keys:
                target_key = key[len("model.") :]
                if target_key not in state_dict:
                    continue
                tensor = shard_tensors[key]
                if tuple(tensor.shape) != tuple(state_dict[target_key].shape):
                    raise ValueError(
                        f"Shape mismatch while loading Qwen3 weight {key}: "
                        f"checkpoint={tuple(tensor.shape)} model={tuple(state_dict[target_key].shape)}"
                    )
                loaded_state[target_key] = tensor

        missing, unexpected = self.load_state_dict(loaded_state, strict=False)
        unexpected = [key for key in unexpected if key]
        if unexpected:
            raise RuntimeError(f"Unexpected Qwen3 pretrained keys after load: {unexpected[:10]}")
        missing = [key for key in missing if key not in loaded_state]
        if missing:
            raise RuntimeError(f"Missing Qwen3 pretrained keys after load: {missing[:10]}")

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[tuple[tuple[torch.Tensor, torch.Tensor], ...]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: bool = True,
        cu_seqlens: Optional[torch.Tensor] = None,
        num_sequences: Optional[torch.Tensor] = None,
    ) -> BaseModelOutputWithPast:
        del input_ids, output_attentions
        if inputs_embeds is None:
            raise ValueError("inputs_embeds must be provided.")

        use_cache = bool(use_cache)
        if use_cache and cu_seqlens is not None:
            raise ValueError("use_cache=True is not supported together with cu_seqlens packing.")

        hidden_states = inputs_embeds
        if attention_mask is None:
            attention_mask = torch.ones(hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device)
        else:
            attention_mask = attention_mask.to(dtype=torch.bool, device=hidden_states.device)
        query_attention_mask = attention_mask[:, -hidden_states.shape[1] :]

        packed_metadata = None
        if position_ids is None:
            if cu_seqlens is not None:
                position_ids = MossTTSNanoGPT2Model.build_packed_position_ids(
                    attention_mask=attention_mask,
                    cu_seqlens=cu_seqlens.to(device=hidden_states.device),
                    num_sequences=num_sequences.to(device=hidden_states.device) if num_sequences is not None else None,
                    sequence_length=hidden_states.shape[1],
                )
            elif attention_mask is not None:
                position_ids = attention_mask.long().cumsum(dim=-1) - 1
                position_ids = position_ids.masked_fill(~attention_mask, 0)
                position_ids = position_ids[:, -hidden_states.shape[1] :]
            else:
                past_length = 0
                if past_key_values is not None and len(past_key_values) > 0:
                    past_length = past_key_values[0][0].shape[1]
                position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device, dtype=torch.long)
                position_ids = position_ids + past_length
                position_ids = position_ids.unsqueeze(0).expand(hidden_states.shape[0], -1)

        if cu_seqlens is not None and self.attn_implementation == "flash_attention_2":
            packed_metadata = MossTTSNanoGPT2Model.build_packed_metadata(
                hidden_states=hidden_states,
                cu_seqlens=cu_seqlens.to(device=hidden_states.device),
                num_sequences=num_sequences.to(device=hidden_states.device) if num_sequences is not None else None,
            )

        hidden_states = hidden_states * query_attention_mask.unsqueeze(-1).to(dtype=hidden_states.dtype)
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        all_hidden_states = () if output_hidden_states else None
        presents = [] if use_cache else None
        for layer_index, decoder_layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            if self.gradient_checkpointing and self.training:
                if use_cache:
                    raise ValueError("use_cache=True is not supported when gradient checkpointing is enabled during training.")

                def custom_forward(*inputs):
                    output, _ = decoder_layer(
                        hidden_states=inputs[0],
                        attention_mask=inputs[1],
                        packed_metadata=packed_metadata,
                        layer_past=None,
                        use_cache=False,
                        position_embeddings=position_embeddings,
                    )
                    return output

                hidden_states = torch.utils.checkpoint.checkpoint(
                    custom_forward,
                    hidden_states,
                    attention_mask,
                    use_reentrant=self.gradient_checkpointing_use_reentrant,
                )
                present = None
            else:
                hidden_states, present = decoder_layer(
                    hidden_states=hidden_states,
                    attention_mask=attention_mask,
                    packed_metadata=packed_metadata,
                    layer_past=None if past_key_values is None else past_key_values[layer_index],
                    use_cache=use_cache,
                    position_embeddings=position_embeddings,
                )

            hidden_states = hidden_states * query_attention_mask.unsqueeze(-1).to(dtype=hidden_states.dtype)
            if presents is not None:
                presents.append(present)

        hidden_states = self.norm(hidden_states)
        hidden_states = hidden_states * query_attention_mask.unsqueeze(-1).to(dtype=hidden_states.dtype)
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return (hidden_states, tuple(presents) if presents is not None else None, all_hidden_states, None)

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=tuple(presents) if presents is not None else None,
            hidden_states=all_hidden_states,
            attentions=None,
        )