Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2025 Jingze Shi and the HuggingFace Inc. team. All rights reserved.3#4# The Doge family of small language models is trained by SmallDoge Team.5#6# Licensed under the Apache License, Version 2.0 (the "License");7# you may not use this file except in compliance with the License.8# You may obtain a copy of the License at9#10#     http://www.apache.org/licenses/LICENSE-2.011#12# Unless required by applicable law or agreed to in writing, software13# distributed under the License is distributed on an "AS IS" BASIS,14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.15# See the License for the specific language governing permissions and16# limitations under the License.17"""PyTorch Doge model."""18 19import math20from typing import Callable, Optional, Union21 22import torch23import torch.nn.functional as F24from torch import nn25 26from ...activations import ACT2FN27from ...cache_utils import Cache28from ...configuration_utils import PretrainedConfig29from ...integrations.flex_attention import compile_friendly_flex_attention30from ...modeling_layers import GradientCheckpointingLayer31from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast32from ...modeling_rope_utils import rope_config_validation33from ...modeling_utils import AttentionInterface, PreTrainedModel34from ...processing_utils import Unpack35from ...utils import TransformersKwargs, is_torch_flex_attn_available36from ...utils.deprecation import deprecate_kwarg37from ...utils.generic import OutputRecorder38from ..llama.modeling_llama import (39    LlamaForSequenceClassification,40    LlamaMLP,41    LlamaPreTrainedModel,42    LlamaRMSNorm,43    LlamaRotaryEmbedding,44    apply_rotary_pos_emb,45    eager_attention_forward,46    repeat_kv,47)48from ..mixtral.modeling_mixtral import MixtralForCausalLM, MixtralModel49 50 51if is_torch_flex_attn_available():52    from torch.nn.attention.flex_attention import BlockMask53 54 55class DogeConfig(PretrainedConfig):56    r"""57    This is the configuration class to store the configuration of a [`DogeModel`]. It is used to instantiate an Doge58    model according to the specified arguments, defining the model architecture like [SmallDoge/Doge-320M](https://huggingface.co/SmallDoge/Doge-320M).59 60    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the61    documentation from [`PretrainedConfig`] for more information.62 63    Args:64        vocab_size (`int`, *optional*, defaults to 32768):65            Vocabulary size of the Doge2 model. Defines the number of different tokens that can be represented by the `inputs_ids` passed when calling [`DogeModel`]66        hidden_size (`int`, *optional*, defaults to 1024):67            Dimension of the hidden representations.68        intermediate_size (`int`, *optional*, defaults to 2048):69            Dimension of the MLP representations.70        num_hidden_layers (`int`, *optional*, defaults to 32):71            Number of hidden layers in the Transformer decoder.72        hidden_dropout (`float`, *optional*, defaults to 0.0):73            Dropout probability for each sequence transformation and state transformation module.74        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):75            The non-linear activation function (function or string) in the decoder.76        initializer_range (`float`, *optional*, defaults to 0.02):77            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.78        rms_norm_eps (`float`, *optional*, defaults to 1e-06):79            The epsilon used by the rms normalization layers.80        use_cache (`bool`, *optional*, defaults to `True`):81            Whether or not the model should return the last key/values attentions (not used by all models). Only82            relevant if `config.is_decoder=True`.83        tie_word_embeddings (`bool`, *optional*, defaults to `False`):84            Whether the model's input and output word embeddings should be tied.85        max_position_embeddings (`int`, *optional*, defaults to 2048):86            The maximum sequence length that this model might ever be used with.87        rope_theta (`float`, *optional*, defaults to 10000.0):88            The base period of the RoPE embeddings.89        rope_scaling (`Dict`, *optional*):90            Dictionary containing the scaling configuration for the RoPE embeddings.91            NOTE: if you apply new rope type and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value accordingly.92            Doge family of small models use `{ 'rope_type': 'dynamic', 'factor': 4.0, 'original_max_position_embeddings': 2048 }` as the default value.93            Expected contents:94                `rope_type` (`str`):95                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope', 'llama3'], with 'default' being the original RoPE implementation.96                `factor` (`float`, *optional*):97                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings.98                    In most scaling types, a `factor` of x will enable the model to handle sequences of length x * original maximum pre-trained length.99                `original_max_position_embeddings` (`int`, *optional*):100                    Used with 'dynamic', 'longrope' and 'llama3'.101                    The original max position embeddings used during pretraining.102                `attention_factor` (`float`, *optional*):103                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention104                    computation.105                    If unspecified, it defaults to value recommended by the implementation, using the `factor` field to infer the suggested value.106                `beta_fast` (`float`, *optional*):107                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear108                    ramp function. If unspecified, it defaults to 32.109                `beta_slow` (`float`, *optional*):110                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear111                    ramp function. If unspecified, it defaults to 1.112                `short_factor` (`List[float]`, *optional*):113                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<`original_max_position_embeddings`).114                    Must be a list of numbers with the same length as the hidden size divided by the number of attention heads divided by 2115                `long_factor` (`List[float]`, *optional*):116                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<`original_max_position_embeddings`).117                    Must be a list of numbers with the same length as the hidden size divided by the number of attention heads divided by 2118                `low_freq_factor` (`float`, *optional*):119                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE120                `high_freq_factor` (`float`, *optional*):121                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE122        num_attention_heads (`int`, *optional*, defaults to 8):123            Number of attention heads for each attention layer in the Transformer decoder.124        num_key_value_heads (`int`, *optional*):125            This is the number of key_value heads that should be used to implement Grouped Query Attention.126            If `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if127            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used.128            When converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed by meanpooling all the original heads within that group.129            For more details checkout [this paper](https://huggingface.co/papers/2305.13245).130            If it is not specified, will default to `num_attention_heads`.131        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):132            Whether to use a bias in the query, key, value and output projection layers during self-attention.133        attention_dropout (`float`, *optional*, defaults to 0.0):134            The dropout ratio for the attention probabilities.135        mlp_bias (`bool`, *optional*, defaults to `False`):136            Whether to use a bias in up_proj, down_proj and gate_proj layers in the MLP layers.137        sliding_window (`int`, *optional*):138            Sliding window attention window size. If not specified, will default to `None`.139        keep_window_size (`int`, *optional*, defaults to 2048):140            The window size of tokens that are not dynamically masked, and dynamic masking is only performed when the sequence length exceeds this value.141        is_moe (`bool`, *optional*, defaults to `False`):142            Whether to use the Cross Domain Mixture of Experts, if `True`, the MoE will inherit the MLP to initialize.143        num_experts (`int`, *optional*, defaults to 16384):144            Number of routed experts in the model. This is only used when `is_moe=True`.145        num_experts_per_tok (`int`, *optional*, defaults to 64):146            Number of selected experts to route per-token.147        norm_topk_prob (`bool`, *optional*, defaults to `False`):148            Whether to normalize the topk probabilities.149        output_router_logits (`bool`, *optional*, defaults to `False`):150            Whether or not the router logits should be returned by the model. Enabling this will also151            allow the model to output the auxiliary loss, including load balancing loss and router z-loss.152        router_aux_loss_coef (`float`, *optional*, defaults to 0.001):153            The aux loss factor for the total loss.154 155    ```python156    >>> from transformers import DogeConfig, DogeModel157 158    >>> # Initializing a Doge-320M style configuration159    >>> configuration = DogeConfig()160 161    >>> # Initializing a model from the Doge-320M style configuration162    >>> model = DogeModel(configuration)163 164    >>> # Accessing the model configuration165    >>> configuration = model.config166    ```"""167 168    model_type = "doge"169    keys_to_ignore_at_inference = ["past_key_values"]170    # Default tensor parallel plan for base model `DogeModel`171    base_model_tp_plan = {172        "layers.*.self_attn.q_proj": "colwise",173        "layers.*.self_attn.k_proj": "colwise",174        "layers.*.self_attn.v_proj": "colwise",175        "layers.*.self_attn.dt_proj": "rowwise",176        "layers.*.self_attn.o_proj": "rowwise",177        "layers.*.input_layernorm.weight": "sequence_parallel",178        "layers.*.input_residual.weight": "sequence_parallel",179        "layers.*.post_attention_layernorm.weight": "sequence_parallel",180        "layers.*.post_attention_residual.weight": "sequence_parallel",181        "norm.weight": "sequence_parallel",182        "layers.*.mlp.gate_proj": "colwise",183        "layers.*.mlp.up_proj": "colwise",184        "layers.*.mlp.down_proj": "rowwise",185        "layers.*.mlp.router_gate": "colwise_rep",186        "layers.*.mlp.down_embed": "rowwise_rep",187        "layers.*.mlp.up_embed": "rowwise_rep",188    }189    base_model_pp_plan = {190        "embed_tokens": (["input_ids"], ["inputs_embeds"]),191        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),192        "norm": (["hidden_states"], ["hidden_states"]),193    }194 195    def __init__(196        self,197        vocab_size=32768,198        hidden_size=1024,199        intermediate_size=2048,200        num_hidden_layers=32,201        hidden_dropout=0.0,202        hidden_act="silu",203        initializer_range=0.02,204        rms_norm_eps=1e-06,205        use_cache=True,206        tie_word_embeddings=False,207        max_position_embeddings=2048,208        rope_theta=10000.0,209        rope_scaling=None,210        num_attention_heads=8,211        num_key_value_heads=None,212        attention_bias=False,213        attention_dropout=0.0,214        mlp_bias=False,215        sliding_window=None,216        keep_window_size=2048,217        is_moe=False,218        num_experts=16384,219        num_experts_per_tok=64,220        norm_topk_prob=False,221        output_router_logits=False,222        router_aux_loss_coef=0.001,223        **kwargs,224    ):225        self.vocab_size = vocab_size226        self.hidden_size = hidden_size227        self.intermediate_size = intermediate_size228        self.num_hidden_layers = num_hidden_layers229 230        self.hidden_dropout = hidden_dropout231        self.hidden_act = hidden_act232        self.initializer_range = initializer_range233        self.rms_norm_eps = rms_norm_eps234        self.use_cache = use_cache235 236        self.max_position_embeddings = max_position_embeddings237        self.rope_theta = rope_theta238        self.rope_scaling = rope_scaling239        self.num_attention_heads = num_attention_heads240        self.num_key_value_heads = num_key_value_heads241        self.attention_bias = attention_bias242        self.attention_dropout = attention_dropout243        self.mlp_bias = mlp_bias244        self.sliding_window = sliding_window245        self.keep_window_size = keep_window_size246        self.is_moe = is_moe247        self.num_experts = num_experts248        self.num_experts_per_tok = num_experts_per_tok249        self.norm_topk_prob = norm_topk_prob250        self.output_router_logits = output_router_logits251        self.router_aux_loss_coef = router_aux_loss_coef252 253        # Validate the correctness of rotary position embeddings parameters254        # BC: if there is a 'type' field, copy it it to 'rope_type'.255        if self.rope_scaling is not None and "type" in self.rope_scaling:256            self.rope_scaling["rope_type"] = self.rope_scaling["type"]257        rope_config_validation(self)258 259        # for backward compatibility260        if num_key_value_heads is None:261            self.num_key_value_heads = num_attention_heads262 263        super().__init__(264            tie_word_embeddings=tie_word_embeddings,265            **kwargs,266        )267 268 269class DogeRMSNorm(LlamaRMSNorm):270    pass271 272 273class DogeRotaryEmbedding(LlamaRotaryEmbedding):274    pass275 276 277def flex_attention_forward(278    module: nn.Module,279    query: torch.Tensor,280    key: torch.Tensor,281    value: torch.Tensor,282    attention_mask: Union[torch.Tensor, "BlockMask"],283    scaling: Optional[float] = None,284    softcap: Optional[float] = None,285    head_mask: Optional[torch.Tensor] = None,286    **kwargs,287) -> tuple[torch.Tensor, torch.Tensor]:288    block_mask = None289    causal_mask = None290    if isinstance(attention_mask, BlockMask):291        block_mask = attention_mask292    else:293        causal_mask = attention_mask294 295    if causal_mask is not None:296        causal_mask = causal_mask[:, :, :, : key.shape[-2]]297 298    def score_mod(score, batch_idx, head_idx, q_idx, kv_idx):299        if softcap is not None:300            score = softcap * torch.tanh(score / softcap)301        if causal_mask is not None:302            score = score + causal_mask[batch_idx][head_idx][q_idx][kv_idx]303        if head_mask is not None:304            score = score + head_mask[batch_idx][head_idx][0][0]305        return score306 307    attn_output, attention_weights = compile_friendly_flex_attention(308        query,309        key,310        value,311        score_mod=score_mod,312        block_mask=block_mask,313        enable_gqa=True,314        scale=scaling,315        # Last time checked on PyTorch == 2.5.1: Flex Attention always computes the lse regardless.316        # For simplification, we thus always return it as no additional computations are introduced.317        return_lse=True,318    )319    # lse is returned in float32320    attention_weights = attention_weights.to(value.dtype)321    attn_output = attn_output.transpose(1, 2).contiguous()322 323    return attn_output, attention_weights324 325 326ALL_ATTENTION_FUNCTIONS = AttentionInterface()327ALL_ATTENTION_FUNCTIONS["doge_flex_attention"] = flex_attention_forward328 329 330class DogeAttention(nn.Module):331    def __init__(self, config: DogeConfig, layer_idx: Optional[int] = None):332        super().__init__()333        self.config = config334        self.layer_idx = layer_idx335        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)336        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads337        self.scaling = self.head_dim**-0.5338        self.attention_dropout = config.attention_dropout339        self.keep_window_size = config.keep_window_size340 341        self.q_proj = nn.Linear(342            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias343        )344        self.k_proj = nn.Linear(345            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias346        )347        self.v_proj = nn.Linear(348            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias349        )350        # dynamic mask for the QK^T attention weights matrix351        self.A = nn.Parameter(torch.zeros(config.num_key_value_heads))352        self.dt_proj = nn.Linear(353            config.num_key_value_heads * self.head_dim, config.num_key_value_heads, bias=config.attention_bias354        )355        self.o_proj = nn.Linear(356            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias357        )358        self.q_norm = DogeRMSNorm(self.head_dim, eps=config.rms_norm_eps)359        self.k_norm = DogeRMSNorm(self.head_dim, eps=config.rms_norm_eps)360 361    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")362    def forward(363        self,364        hidden_states: torch.Tensor,365        position_embeddings: tuple[torch.Tensor, torch.Tensor],366        attention_mask: Optional[torch.Tensor] = None,367        past_key_values: Optional[Cache] = None,368        cache_position: Optional[torch.LongTensor] = None,369        **kwargs,370    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:371        input_shape = hidden_states.shape[:-1]372        hidden_shape = (*input_shape, -1, self.head_dim)373 374        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)375        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)376        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)377 378        cos, sin = position_embeddings379        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)380 381        if past_key_values is not None:382            # sin and cos are specific to RoPE models; cache_position needed for the static cache383            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}384            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)385 386        # calculate dynamic mask from value_states387        dt_states = self.dt_proj(388            value_states.transpose(1, 2).reshape(value_states.shape[0], value_states.shape[-2], -1)389        )390        dt_states = torch.exp(self.A * F.softplus(dt_states)).transpose(-1, -2)391        attn_mask = self.prepare_dynamic_mask(392            hidden_states=hidden_states,393            dt_states=dt_states,394            keep_window_size=self.keep_window_size,395            attention_mask=attention_mask,396        )397        attn_mask = repeat_kv(attn_mask, self.num_key_value_groups)398 399        attention_interface: Callable = eager_attention_forward400        if self.config._attn_implementation != "eager":401            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]402 403        attn_output, attn_weights = attention_interface(404            self,405            query_states,406            key_states,407            value_states,408            attention_mask=attn_mask,409            dropout=0.0 if not self.training else self.attention_dropout,410            scaling=self.scaling,411            **kwargs,412        )413 414        attn_output = attn_output.reshape(*input_shape, -1).contiguous()415        attn_output = self.o_proj(attn_output)416        return attn_output, attn_weights417 418    def prepare_dynamic_mask(419        self,420        hidden_states: torch.Tensor,421        dt_states: torch.Tensor,422        keep_window_size: int = 2048,423        attention_mask: Optional[torch.Tensor] = None,424    ):425        """426        The core idea of DMA is to calculate the dynamic attention mask to mask the tokens that should be masked, so as to form sparse attention.427 428        Combine `dt_states` with `attention_mask` to generate the final `attn_mask`.429 430        Args:431            hidden_states (`torch.Tensor`): The input hidden_states, used to determine the minimum value of the current input precision.432            dt_states (`torch.Tensor`): dt_states of shape `(batch_size, num_heads, key_sequence_length)`.433            keep_window_size (`int`): The window size of tokens that are not dynamically masked, and dynamic masking is only performed when the sequence length exceeds this value.434            attention_mask (`torch.Tensor`, *optional*): attention mask of shape `(batch_size, 1, query_sequence_length, key_sequence_length)`.435        """436        min_dtype = torch.finfo(hidden_states.dtype).min437        dtype = hidden_states.dtype438        attn_mask = dt_states[:, :, None, :].expand(439            -1, -1, hidden_states.shape[1], -1440        )  # [batch_size, num_heads, query_len, key_len]441        if attention_mask is not None and not isinstance(attention_mask, BlockMask):442            if attention_mask.dtype == torch.bool:443                dtype = hidden_states.dtype444                attention_mask = torch.where(445                    attention_mask, torch.tensor(0.0, device=attention_mask.device, dtype=dtype), min_dtype446                )447            attn_mask = attn_mask.masked_fill(attention_mask[:, :, :, : attn_mask.shape[-1]] != 0, min_dtype)448        if attn_mask.shape[-1] > keep_window_size:449            active_mask = torch.zeros_like(attn_mask, dtype=dtype, device=attn_mask.device)450            topk_indices = torch.topk(attn_mask, keep_window_size, dim=-1, largest=True, sorted=False).indices451            active_mask = active_mask.scatter(-1, topk_indices, 1.0)452            attn_mask = attn_mask.masked_fill(active_mask == 0.0, min_dtype)453        return attn_mask454 455 456class DogeMLP(LlamaMLP):457    pass458 459 460class DogeCDMoE(nn.Module):461    def __init__(self, config: DogeConfig):462        super().__init__()463        self.hidden_size = config.hidden_size464        self.intermediate_size = config.intermediate_size465        self.act_fn = ACT2FN[config.hidden_act]466 467        self.num_experts = config.num_experts468        self.num_keys = math.floor(math.sqrt(self.num_experts))469        self.top_k = config.num_experts_per_tok470        self.norm_topk_prob = config.norm_topk_prob471 472        # shared expert473        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)474        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)475        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)476 477        # router gate for retrieval experts478        self.router_gate = nn.Linear(self.hidden_size, self.num_keys * 2, bias=False)479 480        # routed experts481        self.down_embed = nn.Embedding(self.num_experts, self.hidden_size)482        self.up_embed = nn.Embedding(self.num_experts, self.hidden_size)483 484    def forward(485        self,486        hidden_states: torch.Tensor,487        **kwargs,488    ) -> torch.Tensor:489        bsz, seq_len, _ = hidden_states.shape490 491        # get routing logits with router gate492        router_logits = self.router_gate(hidden_states).view(2, bsz * seq_len, -1)493 494        # get experts with the highest routing logits495        (scores_x, scores_y), (indices_x, indices_y) = router_logits.topk(self.num_keys, dim=-1)496        all_scores = scores_x.unsqueeze(-1) + scores_y.unsqueeze(-2)497        all_indices = indices_x.unsqueeze(-1) * self.num_keys + indices_y.unsqueeze(-2)498        all_scores = all_scores.view(*all_scores.shape[:-2], -1)499        all_indices = all_indices.view(*all_indices.shape[:-2], -1)500        scores, position_indices = all_scores.topk(self.top_k, dim=-1)501        indices = all_indices.gather(-1, position_indices)502        routing_weights = F.softmax(scores, dim=-1)503        if self.norm_topk_prob:504            routing_weights /= routing_weights.sum(dim=-1, keepdim=True)505 506        # mix routed experts states with shared expert states507        down_embed = self.down_embed(indices)508        up_embed = self.up_embed(indices)509        experts_weights = torch.matmul(down_embed, hidden_states.view(bsz * seq_len, -1, 1)).view(bsz * seq_len, -1)510        experts_weights = self.act_fn(experts_weights) * routing_weights511        experts_states = torch.matmul(experts_weights.view(bsz * seq_len, 1, -1), up_embed).view(bsz, seq_len, -1)512        hidden_states = self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))513        hidden_states = hidden_states + experts_states514        return hidden_states, router_logits515 516 517class DogeDecoderLayer(GradientCheckpointingLayer):518    def __init__(self, config: DogeConfig, layer_idx: Optional[int] = None):519        super().__init__()520        self.hidden_dropout = config.hidden_dropout521 522        self.input_layernorm = DogeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)523        self.self_attn = DogeAttention(config=config, layer_idx=layer_idx)524        self.input_residual = nn.Parameter(torch.ones(config.hidden_size))525 526        self.post_attention_layernorm = DogeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)527        self.mlp = DogeMLP(config) if not config.is_moe else DogeCDMoE(config)528        self.post_attention_residual = nn.Parameter(torch.ones(config.hidden_size))529 530    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")531    def forward(532        self,533        hidden_states: torch.Tensor,534        position_embeddings: tuple[torch.Tensor, torch.Tensor],535        attention_mask: Optional[torch.Tensor] = None,536        position_ids: Optional[torch.LongTensor] = None,537        past_key_values: Optional[Cache] = None,538        use_cache: Optional[bool] = False,539        cache_position: Optional[torch.LongTensor] = None,540        **kwargs: Unpack[TransformersKwargs],541    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:542        # sequence transformation543        residual = hidden_states544        hidden_states = self.input_layernorm(hidden_states)545        hidden_states, self_attn_weights = self.self_attn(546            hidden_states=hidden_states,547            position_embeddings=position_embeddings,548            attention_mask=attention_mask,549            position_ids=position_ids,550            past_key_values=past_key_values,551            use_cache=use_cache,552            cache_position=cache_position,553            **kwargs,554        )555        hidden_states = F.dropout(hidden_states, p=self.hidden_dropout, training=self.training)556        hidden_states = self.input_residual * residual + hidden_states557 558        # state transformation559        residual = hidden_states560        hidden_states = self.post_attention_layernorm(hidden_states)561        hidden_states = self.mlp(hidden_states)562        hidden_states = F.dropout(hidden_states, p=self.hidden_dropout, training=self.training)563        hidden_states = self.post_attention_residual * residual + hidden_states564 565        return hidden_states566 567 568class DogePreTrainedModel(LlamaPreTrainedModel):569    _supports_flash_attn = False570    _can_compile_fullgraph = False571    _can_record_outputs = {572        "router_logits": OutputRecorder(DogeCDMoE, index=1),573        "hidden_states": DogeDecoderLayer,574        "attentions": DogeAttention,575    }576 577    def _init_weights(self, module):578        """Initialize the weights"""579        PreTrainedModel._init_weights(self, module)580        if isinstance(module, DogeAttention):581            if hasattr(module, "A"):582                module.A.data.zero_()583        elif isinstance(module, DogeDecoderLayer):584            if hasattr(module, "input_residual"):585                module.input_residual.data.fill_(1.0)586            if hasattr(module, "post_attention_residual"):587                module.post_attention_residual.data.fill_(1.0)588 589 590class DogeModel(MixtralModel):591    pass592 593 594def load_balancing_loss_func(595    gate_logits: Union[torch.Tensor, tuple[torch.Tensor], None],596    num_experts: Optional[int] = None,597    num_keys: Optional[int] = None,598    top_k: int = 2,599    attention_mask: Optional[torch.Tensor] = None,600) -> Union[torch.Tensor, int]:601    r"""602    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.603 604    See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss605    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between606    experts is too unbalanced.607 608    Args:609        gate_logits:610            Logits from the `router_gate`, should be a tuple of model.config.num_hidden_layers tensors of611            shape [2, batch_size * sequence_length, num_keys].612        num_experts:613            Number of experts614        num_keys:615            Number of keys616        top_k:617            The number of experts to route per-token, can be also interpreted as the `top-k` routing618            parameter.619        attention_mask (`torch.Tensor`, *optional*):620            The attention_mask used in forward function621            shape [batch_size X sequence_length] if not None.622 623    Returns:624        The auxiliary loss.625    """626    if gate_logits is None or not isinstance(gate_logits, tuple):627        return 0628 629    compute_dtype = gate_logits[0].dtype630    compute_device = gate_logits[0].device631    all_expert_indices = []632    all_routing_weights = []633 634    for layer_gate_logits in gate_logits:635        layer_gate_logits = layer_gate_logits.to(compute_device)636 637        (scores_x, scores_y), (indices_x, indices_y) = layer_gate_logits.topk(num_keys, dim=-1)638 639        all_scores = scores_x.unsqueeze(-1) + scores_y.unsqueeze(-2)640        all_indices = indices_x.unsqueeze(-1) * num_keys + indices_y.unsqueeze(-2)641        all_scores = all_scores.view(*all_scores.shape[:-2], -1)642        all_indices = all_indices.view(*all_indices.shape[:-2], -1)643 644        _, position_indices = all_scores.topk(top_k, dim=-1)645        expert_indices = all_indices.gather(-1, position_indices)646 647        routing_weights = F.softmax(all_scores, dim=-1)648 649        all_expert_indices.append(expert_indices)650        all_routing_weights.append(routing_weights)651    all_expert_indices = torch.cat(all_expert_indices, dim=0)652    all_routing_weights = torch.cat(all_routing_weights, dim=0)653 654    if attention_mask is None:655        # Compute the percentage of tokens routed to each experts656        all_expert_indices = all_expert_indices.view(-1)657        tokens_per_expert = torch.zeros(num_experts, dtype=compute_dtype, device=compute_device)658        pad = torch.ones_like(all_expert_indices, dtype=compute_dtype, device=compute_device)659        tokens_per_expert = tokens_per_expert.scatter_add_(0, all_expert_indices, pad) / all_expert_indices.shape[0]660 661        # Compute the average probability of routing to these experts662        router_prob_per_expert = torch.mean(all_routing_weights, dim=0)663    else:664        batch_size, sequence_length = attention_mask.shape665        num_hidden_layers = len(gate_logits)666 667        #  Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask668        expert_attention_mask = (669            attention_mask[None, :, :, None]670            .expand((num_hidden_layers, batch_size, sequence_length, top_k))671            .reshape(-1)672            .to(compute_device)673        )674        all_expert_indices = all_expert_indices.view(-1)[expert_attention_mask.bool()]675 676        # Compute the percentage of tokens routed to each experts677        tokens_per_expert = torch.zeros(num_experts, dtype=compute_dtype, device=compute_device)678        pad = torch.ones_like(all_expert_indices, dtype=compute_dtype, device=compute_device)679        tokens_per_expert = tokens_per_expert.scatter_add_(0, all_expert_indices, pad) / torch.sum(680            expert_attention_mask681        )682 683        # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert684        router_per_expert_attention_mask = (685            attention_mask[None, :, :, None]686            .expand((num_hidden_layers, batch_size, sequence_length, num_experts))687            .reshape(-1, num_experts)688            .to(compute_device)689        )690 691        # Compute the average probability of routing to these experts692        router_prob_per_expert = torch.sum(all_routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(693            router_per_expert_attention_mask, dim=0694        )695 696    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert)697    return overall_loss * num_experts698 699 700class DogeForCausalLM(MixtralForCausalLM):701    def __init__(self, config):702        super().__init__(config)703        self.model = DogeModel(config)704        self.num_experts = config.num_experts705 706    def forward(707        self,708        input_ids: Optional[torch.LongTensor] = None,709        attention_mask: Optional[torch.Tensor] = None,710        position_ids: Optional[torch.LongTensor] = None,711        past_key_values: Optional[Cache] = None,712        inputs_embeds: Optional[torch.FloatTensor] = None,713        labels: Optional[torch.LongTensor] = None,714        use_cache: Optional[bool] = None,715        cache_position: Optional[torch.LongTensor] = None,716        logits_to_keep: Union[int, torch.Tensor] = 0,717        output_router_logits: Optional[bool] = None,718        **kwargs: Unpack[TransformersKwargs],719    ) -> MoeCausalLMOutputWithPast:720        r"""721        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):722            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,723            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored724            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.725 726        Example:727 728        ```python729        >>> from transformers import AutoTokenizer, DogeForCausalLM730 731        >>> model = DogeForCausalLM.from_pretrained("SmallDoge/Doge-320M")732        >>> tokenizer = AutoTokenizer.from_pretrained("SmallDoge/Doge-320M")733 734        >>> prompt = "Hey, are you conscious? Can you talk to me?"735        >>> inputs = tokenizer(prompt, return_tensors="pt")736 737        >>> # Generate738        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)739        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]740        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."741        ```"""742        output_router_logits = (743            output_router_logits if output_router_logits is not None else self.config.output_router_logits744        )745 746        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)747        outputs: MoeModelOutputWithPast = self.model(748            input_ids=input_ids,749            attention_mask=attention_mask,750            position_ids=position_ids,751            past_key_values=past_key_values,752            inputs_embeds=inputs_embeds,753            use_cache=use_cache,754            cache_position=cache_position,755            **kwargs,756        )757 758        hidden_states = outputs.last_hidden_state759        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss760        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep761        logits = self.lm_head(hidden_states[:, slice_indices, :])762 763        loss = None764        if labels is not None:765            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)766 767        aux_loss = None768        if output_router_logits:769            aux_loss = load_balancing_loss_func(770                outputs.router_logits,771                self.num_experts,772                math.floor(math.sqrt(self.num_experts)),773                self.num_experts_per_tok,774                attention_mask,775            )776            if labels is not None:777                loss += self.router_aux_loss_coef * aux_loss.to(loss.device)  # make sure to reside in the same device778 779        return MoeCausalLMOutputWithPast(780            loss=loss,781            aux_loss=aux_loss,782            logits=logits,783            past_key_values=outputs.past_key_values,784            hidden_states=outputs.hidden_states,785            attentions=outputs.attentions,786            router_logits=outputs.router_logits,787        )788 789 790class DogeForSequenceClassification(LlamaForSequenceClassification):791    pass792 793 794__all__ = [795    "DogeConfig",796    "DogeForCausalLM",797    "DogeModel",798    "DogePreTrainedModel",799    "DogeForSequenceClassification",800]801