baichuan_moe.py

# 2024 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
import time
"""Inference-only Baichuan-MOE model."""
from transformers.configuration_utils import PretrainedConfig
class BaiChuanMoEConfig(PretrainedConfig):
    model_type = "baichuan-moe"
    keys_to_ignore_at_inference = ["past_key_values"]

    def __init__(
        self,
        vocab_size=64000,
        hidden_size=4096,
        intermediate_size=11008,
        num_hidden_layers=32,
        num_attention_heads=32,
        hidden_act="silu",
        max_position_embeddings=4096,
        initializer_range=0.02,
        rms_norm_eps=1e-6,
        rope_base=1e6,
        use_cache=True,
        pad_token_id=0,
        bos_token_id=1,
        eos_token_id=2,
        tie_word_embeddings=False,
        moe_experts_fixed=0,
        moe_experts_selected=2,
        moe_experts_routed=8,
        num_experts_fixed_per_layer=None, # "0,0,0,1,0,2"
        num_experts_selected_per_layer=None, # "1,2,1,1,1,2"
        num_experts_routed_per_layer=None, # "1,8,1,8,1,16"
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.max_position_embeddings = max_position_embeddings
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.hidden_act = hidden_act
        self.initializer_range = initializer_range
        self.rms_norm_eps = rms_norm_eps
        self.rope_base = rope_base
        self.use_cache = use_cache
        self.moe_experts_fixed = moe_experts_fixed
        self.moe_experts_selected = moe_experts_selected
        self.moe_experts_routed = moe_experts_routed
        if num_experts_routed_per_layer:
            self.num_experts_routed_per_layer = [int(_.strip()) for _ in num_experts_routed_per_layer.split(",")]
            assert len(self.num_experts_routed_per_layer) == self.num_hidden_layers
            assert all([_ >= 1 for _ in self.num_experts_routed_per_layer])
        else:
            self.num_experts_routed_per_layer = None

        if num_experts_selected_per_layer:
            self.num_experts_selected_per_layer = [int(_.strip()) for _ in num_experts_selected_per_layer.split(",")]
            assert len(self.num_experts_selected_per_layer) == self.num_hidden_layers
            assert all([x >= y for x, y in zip(self.num_experts_routed_per_layer, self.num_experts_selected_per_layer)])
        else:
            self.num_experts_selected_per_layer = None

        if num_experts_fixed_per_layer:
            self.num_experts_fixed_per_layer = [int(_.strip()) for _ in num_experts_fixed_per_layer.split(",")]
            assert len(self.num_experts_fixed_per_layer) == self.num_hidden_layers
        else:
            self.num_experts_fixed_per_layer = None

        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )
import copy
import math
from typing import Iterable, List, Optional, Tuple, Union

import torch
from torch import nn
from transformers import PretrainedConfig

from vllm.attention import Attention, AttentionMetadata
from vllm.config import CacheConfig, LoRAConfig
from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
                              get_tensor_model_parallel_world_size,
                              tensor_model_parallel_all_reduce)
from vllm.model_executor.layers.fused_moe import fused_moe


from vllm.model_executor.layers.activation import SiluAndMul,GeluAndMul
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
                                               QKVParallelLinear,
                                               ReplicatedLinear,
                                               RowParallelLinear,
                                               LinearMethodBase)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
from vllm.model_executor.layers.vocab_parallel_embedding import (
    DEFAULT_VOCAB_PADDING_SIZE,ParallelLMHead, VocabParallelEmbedding)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader, maybe_remap_kv_scale_name
from vllm.model_executor.utils import set_weight_attrs
from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.sequence import IntermediateTensors

from .interfaces import SupportsLoRA, SupportsPP
from .utils import (is_pp_missing_parameter, PPMissingLayer,
                    make_empty_intermediate_tensors_factory, make_layers)

def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
    closest_power_of_2 = 2**math.floor(math.log2(total_num_heads))
    base = torch.tensor(
        2**(-(2**-(math.log2(closest_power_of_2) - 3))),
        dtype=torch.float32,
    )
    powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32)
    slopes = torch.pow(base, powers)

    if closest_power_of_2 != total_num_heads:
        extra_base = torch.tensor(
            2**(-(2**-(math.log2(2 * closest_power_of_2) - 3))),
            dtype=torch.float32,
        )
        num_remaining_heads = min(closest_power_of_2,
                                  total_num_heads - closest_power_of_2)
        extra_powers = torch.arange(start=1,
                                    end=1 + 2 * num_remaining_heads,
                                    step=2,
                                    dtype=torch.int32)
        slopes = torch.cat(
            [slopes, torch.pow(extra_base, extra_powers)], dim=0)
    return slopes

class MLP(nn.Module):
    def __init__(
            self,
            hidden_size: int,
            intermediate_size: int,
            hidden_act: str,
            layer_index: int
    ):
        super().__init__()
        self.layer_index = layer_index
        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size, [intermediate_size] * 2,
            bias=False,)
            # linear_method=None)
        self.down_proj = RowParallelLinear(intermediate_size,
                                           hidden_size,
                                           bias=False,)
                                        #    linear_method=None)
        if hidden_act not in ["silu", "gelu"]:
            raise ValueError(f"Unsupported activation: {hidden_act}. "
                             "Only silu and gelu are supported for now.")
        self.act_fn = SiluAndMul() if hidden_act == "silu" else GeluAndMul()

    def forward(self, x):
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        ret, _ = self.down_proj(x)

        return ret


class MixtralMLP(nn.Module):
    """
    This implementation is
    strictly equivalent to standard MoE with full capacity (no
    dropped tokens). It's faster since it formulates MoE operations
    in terms of block-sparse operations to accomodate imbalanced
    assignments of tokens to experts, whereas standard MoE either
    (1) drop tokens at the cost of reduced performance or (2) set
    capacity factor to number of experts and thus waste computation
    and memory on padding.
    """

    def __init__(self,
                hidden_size,
                intermediate_size,
                hidden_act,
                moe_experts_routed,
                moe_experts_selected,
                moe_experts_fixed,
                layer_index,
                params_dtype: Optional[torch.dtype] = None,
                tp_size: Optional[int] = None):
        super().__init__()

        self.layer_index = layer_index
        self.tp_size = tp_size or get_tensor_model_parallel_world_size()

        self.num_experts_routed = moe_experts_routed
        self.num_local_experts_routed = self.num_experts_routed // 1
        self.top_k = moe_experts_selected
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size // self.tp_size


        if params_dtype is None:
            params_dtype = torch.get_default_dtype()
        self.params_dtype = params_dtype
        self.router = ReplicatedLinear(self.hidden_size,
                                        self.num_experts_routed,
                                        bias=False,
                                        params_dtype=self.params_dtype,)
                                        # linear_method=None)

        self.ws = nn.Parameter(
            torch.empty(self.num_experts_routed,
                        2 * self.intermediate_size,
                        self.hidden_size,
                        device="cuda",
                        dtype=self.params_dtype))
        self.w2s = nn.Parameter(
            torch.empty(self.num_experts_routed,
                        self.hidden_size,
                        self.intermediate_size,
                        device="cuda",
                        dtype=self.params_dtype))
        
        set_weight_attrs(self.ws, {
            "weight_loader": self.weight_loader,
        })
        set_weight_attrs(self.w2s, {
            "weight_loader": self.weight_loader,
        })


        if moe_experts_fixed >= 1:
            self.local_experts_fixed = MLP(hidden_size, intermediate_size*moe_experts_fixed, hidden_act, layer_index)
        else:
            self.local_experts_fixed = None

    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor,
                      weight_name: str, expert_id: int):
        tp_rank = get_tensor_model_parallel_rank()
        param_data = param.data
        shard_size = self.intermediate_size
        shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)
        if weight_name.endswith("gate_proj.weight"):
            param_data[expert_id, 0:shard_size, :] = loaded_weight[shard, :]
        if weight_name.endswith("up_proj.weight"):
            param_data[expert_id,
                       shard_size:2 * shard_size, :] = loaded_weight[shard, :]
        if weight_name.endswith("down_proj.weight"):
            param_data[expert_id, :, :] = loaded_weight[:, shard]


    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """ """
        num_tokens, hidden_size = hidden_states.shape
        hidden_states = hidden_states.view(-1, self.hidden_size)
        router_logits, _ = self.router(hidden_states)
        final_hidden_states = fused_moe(hidden_states,
                                        self.ws,
                                        self.w2s,
                                        router_logits,
                                        self.top_k,
                                        renormalize=True)

        if self.tp_size > 1:
            final_hidden_states = tensor_model_parallel_all_reduce(
                final_hidden_states)
        
        final_hidden_states = final_hidden_states.view(num_tokens, hidden_size)
        
        if self.local_experts_fixed:
            final_hidden_states += self.local_experts_fixed(hidden_states).reshape(num_tokens, hidden_size)
            final_hidden_states /= 2
        
        ret = final_hidden_states.reshape(num_tokens, hidden_size)
        return ret


class MixtralAttention(nn.Module):

    def __init__(self,
                 hidden_size: int,
                 num_heads: int,
                 num_kv_heads: int,
                 max_position: int = 4096 * 32,
                 rope_theta: float = 10000,
                 linear_method: Optional[LinearMethodBase] = None,) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        tp_size = get_tensor_model_parallel_world_size()
        self.total_num_heads = num_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = num_kv_heads
        if self.total_num_kv_heads >= tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_kv_heads % tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
        self.head_dim = hidden_size // self.total_num_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.rope_theta = rope_theta

        self.W_pack = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=False,
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
        )
        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position,
            base=int(self.rope_theta),
            is_neox_style=True,
        )
        self.attn = Attention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
    ) -> torch.Tensor:
        qkv, _ = self.W_pack(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(positions, q, k)
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
        output, _ = self.o_proj(attn_output)
        return output


class DecoderLayer(nn.Module):
    def __init__(
        self,
        config: BaiChuanMoEConfig,
        linear_method: Optional[LinearMethodBase] = None,
        layer_index: Optional[int] = 1
    ) -> None:
        super().__init__()

        self.layer_index = layer_index
        self.hidden_size = config.hidden_size
        # Requires transformers > 4.32.0
        rope_theta = getattr(config, "rope_base", 10000)
        self.self_attn = MixtralAttention(
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            max_position=config.max_position_embeddings,
            num_kv_heads=config.num_attention_heads,
            rope_theta=rope_theta,
            linear_method=linear_method)
        
        
        # Dense
        if config.moe_experts_routed == 1:
            self.mlp = MLP(hidden_size=config.hidden_size,
                            intermediate_size=config.intermediate_size,
                            hidden_act=config.hidden_act, layer_index=layer_index)
        # MoE
        else:
            self.mlp = MixtralMLP(config.hidden_size,
                                    config.intermediate_size,
                                    config.hidden_act,
                                    config.moe_experts_routed,
                                    config.moe_experts_selected,
                                    config.moe_experts_fixed,
                                    layer_index)
        self.input_layernorm = RMSNorm(config.hidden_size,
                                       eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size,
                                                eps=config.rms_norm_eps)
    
    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
        residual: Optional[torch.Tensor],
    ) -> torch.Tensor:

        # Self Attention
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(
                hidden_states, residual)
        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
            kv_cache=kv_cache,
            attn_metadata=attn_metadata,
        )

        # Fully Connected
        hidden_states, residual = self.post_attention_layernorm(
            hidden_states, residual)
        
        hidden_states = self.mlp(hidden_states)

        return hidden_states, residual


class Model(nn.Module):
    def __init__(
        self,
        config: BaiChuanMoEConfig,
        linear_method: Optional[LinearMethodBase] = None,
        lora_config: Optional[LoRAConfig] = None,
    ) -> None:
        super().__init__()
        self.padding_idx = config.pad_token_id
        lora_vocab = (lora_config.lora_extra_vocab_size *
                      (lora_config.max_loras or 1)) if lora_config else 0
        self.vocab_size = config.vocab_size + lora_vocab
        self.org_vocab_size = config.vocab_size

        self.embed_tokens = VocabParallelEmbedding(
            self.vocab_size,
            config.hidden_size,
            org_num_embeddings=config.vocab_size,
        )

        if config.num_experts_routed_per_layer:
            layers = []
            for index in range(config.num_hidden_layers):
                config_ = copy.deepcopy(config)
                config_.moe_experts_fixed = config.num_experts_fixed_per_layer[index]
                config_.moe_experts_selected = config.num_experts_selected_per_layer[index]
                config_.moe_experts_routed = config.num_experts_routed_per_layer[index]

                layers.append(DecoderLayer(config_, linear_method=None, layer_index=index))
            self.layers = nn.ModuleList(layers)
        else:
            self.layers = nn.ModuleList([
                DecoderLayer(config, linear_method=linear_method)
                for _ in range(config.num_hidden_layers)
            ])


        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))
    
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors],
    ) -> torch.Tensor:
        if get_pp_group().is_first_rank:
            hidden_states = self.embed_tokens(input_ids)
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

        hidden_states = self.embed_tokens(input_ids)
        residual = None

        for idx, decoder_layer in enumerate(self.layers):
            hidden_states, residual = decoder_layer(positions,  hidden_states,
                                                    kv_caches[idx], attn_metadata,
                                                    residual)
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({
                "hidden_states": hidden_states,
                "residual": residual
            })        
        hidden_states, _ = self.norm(hidden_states, residual)
        
        return hidden_states

class NormHead(nn.Module):
    def __init__(self, hidden_size, vocab_size, bias=False):
        super().__init__()
        self.weight = nn.Parameter(torch.empty((vocab_size, hidden_size)))
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        self.norm_weight = nn.functional.normalize(self.weight)

    def forward(self, hidden_states):
        return nn.functional.linear(hidden_states, self.norm_weight)

class BaiChuanMoEForCausalLM(nn.Module):
    # packed_modules_mapping = {
    #     "qkv_proj": [
    #         "q_proj",
    #         "k_proj",
    #         "v_proj",
    #     ],
    # }

    # # LoRA specific attributes
    # supported_lora_modules = [
    #     "qkv_proj",
    #     "o_proj",
    #     "embed_tokens",
    #     "lm_head",
    # ]
    embedding_modules = {
        "embed_tokens": "input_embeddings",
        "lm_head": "output_embeddings",
    }
    embedding_padding_modules = ["lm_head"]

    def __init__(
        self,
        config: BaiChuanMoEConfig,
        linear_method: Optional[LinearMethodBase] = None,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
        lora_config: Optional[LoRAConfig] = None,
    ) -> None:
        super().__init__()
        self.config = config
        self.linear_method = linear_method
        self.model = Model(config,
                                  linear_method,
                                  lora_config=lora_config)
        # if get_pp_group().is_last_rank:
        #     self.unpadded_vocab_size = config.vocab_size
        #     if lora_config:
        #         self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
        #     self.lm_head = ParallelLMHead(
        #         self.unpadded_vocab_size,
        #         config.hidden_size,
        #         org_num_embeddings=config.vocab_size,
        #         padding_size=(
        #             DEFAULT_VOCAB_PADDING_SIZE
        #             # We need bigger padding if using lora for kernel
        #             # compatibility
        #             if not lora_config else
        #             lora_config.lora_vocab_padding_size),
        #         quant_config=quant_config,
        #     )
        #     if config.tie_word_embeddings:
        #         self.lm_head = self.lm_head.tie_weights(
        #             self.model.embed_tokens)

        #     logit_scale = getattr(config, "logit_scale", 1.0)
        #     self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
        #                                             config.vocab_size,
        #                                             logit_scale)
        #     self.sampler = Sampler()
        # else:
        #     self.lm_head = PPMissingLayer()
        self.unpadded_vocab_size = config.vocab_size
        if lora_config:
            self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
        self.lm_head = ParallelLMHead(
            self.unpadded_vocab_size,
            config.hidden_size,
            org_num_embeddings=config.vocab_size,
            padding_size=DEFAULT_VOCAB_PADDING_SIZE
            # We need bigger padding if using lora for kernel
            # compatibility
            if not lora_config else lora_config.lora_vocab_padding_size,
            quant_config=quant_config,
        )
        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight
        self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
                                                config.vocab_size)
        self.sampler = Sampler()
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
    
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors] = None,
    ) -> torch.Tensor:
        hidden_states = self.model(input_ids, positions, kv_caches,
                                   attn_metadata, intermediate_tensors)
        return hidden_states

    def compute_logits(self, hidden_states: torch.Tensor,
                       sampling_metadata: SamplingMetadata) -> torch.Tensor:
        logits = self.logits_processor(self.lm_head, hidden_states,
                                       sampling_metadata)
        return logits

    def sample(
        self,
        logits: Optional[torch.Tensor],
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
        next_tokens = self.sampler(logits, sampling_metadata)
        return next_tokens

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
            ("mlp.gate_up_proj", "mlp.gate_proj", 0),
            ("mlp.gate_up_proj", "mlp.up_proj", 1),
            ("mlp.local_experts_fixed.gate_up_proj", "mlp.local_experts_fixed.gate_proj", 0),
            ("mlp.local_experts_fixed.gate_up_proj", "mlp.local_experts_fixed.up_proj", 1),
        ]

        expert_params_mapping = [
            # (param_name, weight_name, expert_id)
            ("ws" if weight_name in ["gate_proj", "up_proj"] else "w2s",
             f"local_experts_routed.{expert_id}.{weight_name}.weight", expert_id)
            for expert_id in range(16)
            for weight_name in ["gate_proj", "down_proj", "up_proj"]
        ]

        params_dict = dict(self.named_parameters())
        
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                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)
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                for param_name, weight_name, expert_id in expert_params_mapping:
                    if weight_name not in name:
                        continue
                    name = name.replace(weight_name, param_name)
                    param = params_dict[name]
                    weight_loader = param.weight_loader
                    weight_loader(param,
                                  loaded_weight,
                                  weight_name,
                                  expert_id=expert_id)
                    break
                else:
                    # Skip loading extra bias for GPTQ models.
                    if name.endswith(".bias") and name not in params_dict:
                        continue

                    param = params_dict.get(name, None)

                    if name == "lm_head.weight":
                        # do norm
                        norm_weight = nn.functional.normalize(loaded_weight)
                        weight_loader = getattr(param, "weight_loader",
                                                default_weight_loader)
                        weight_loader(param, norm_weight)
                    else:
                        weight_loader = getattr(param, "weight_loader",
                                                default_weight_loader)
                        weight_loader(param, loaded_weight)
# 'model.layers.0.mlp.down_proj.weight_packed'