# Adapted from an upstream configuration file. from transformers.configuration_utils import PretrainedConfig from transformers.utils import logging logger = logging.get_logger(__name__) RIZE_PRETRAINED_CONFIG_ARCHIVE_MAP = {} class RizeConfig(PretrainedConfig): r""" This is the configuration class to store the configuration of a [`RizeModel`]. It is used to instantiate an Rize model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the Rize-V3. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. Args: vocab_size (`int`, *optional*, defaults to 129280): Vocabulary size of the Rize model. Defines the number of different tokens that can be represented by the `inputs_ids` passed when calling [`RizeModel`] hidden_size (`int`, *optional*, defaults to 4096): Dimension of the hidden representations. intermediate_size (`int`, *optional*, defaults to 11008): Dimension of the MLP representations. moe_intermediate_size (`int`, *optional*, defaults to 1407): Dimension of the MoE representations. num_hidden_layers (`int`, *optional*, defaults to 32): Number of hidden layers in the Transformer decoder. num_nextn_predict_layers (`int`, *optional*, defaults to 1): Number of nextn predict layers in the RizeV3 Model. num_attention_heads (`int`, *optional*, defaults to 32): Number of attention heads for each attention layer in the Transformer decoder. n_shared_experts (`int`, *optional*, defaults to None): Number of shared experts, None means dense model. n_routed_experts (`int`, *optional*, defaults to None): Number of routed experts, None means dense model. routed_scaling_factor (`float`, *optional*, defaults to 1.0): Scaling factor or routed experts. topk_method (`str`, *optional*, defaults to `gready`): Topk method used in routed gate. n_group (`int`, *optional*, defaults to None): Number of groups for routed experts. topk_group (`int`, *optional*, defaults to None): Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups). num_experts_per_tok (`int`, *optional*, defaults to None): Number of selected experts, None means dense model. moe_layer_freq (`int`, *optional*, defaults to 1): The frequency of the MoE layer: one expert layer for every `moe_layer_freq - 1` dense layers. first_k_dense_replace (`int`, *optional*, defaults to 0): Number of dense layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head). \--k dense layers--/ norm_topk_prob (`bool`, *optional*, defaults to False): Whether to normalize the weights of the routed experts. scoring_func (`str`, *optional*, defaults to 'softmax'): Method of computing expert weights. aux_loss_alpha (`float`, *optional*, defaults to 0.001): Auxiliary loss weight coefficient. seq_aux = (`bool`, *optional*, defaults to True): Whether to compute the auxiliary loss for each individual sample. num_key_value_heads (`int`, *optional*): This is the number of key_value heads that should be used to implement Grouped Query Attention. If `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. 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. For more details checkout [this paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `num_attention_heads`. hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): The non-linear activation function (function or string) in the decoder. max_position_embeddings (`int`, *optional*, defaults to 2048): The maximum sequence length that this model might ever be used with. initializer_range (`float`, *optional*, defaults to 0.02): The standard deviation of the truncated_normal_initializer for initializing all weight matrices. rms_norm_eps (`float`, *optional*, defaults to 1e-06): The epsilon used by the rms normalization layers. use_cache (`bool`, *optional*, defaults to `True`): Whether or not the model should return the last key/values attentions (not used by all models). Only relevant if `config.is_decoder=True`. pad_token_id (`int`, *optional*): Padding token id. bos_token_id (`int`, *optional*, defaults to 1): Beginning of stream token id. eos_token_id (`int`, *optional*, defaults to 2): End of stream token id. pretraining_tp (`int`, *optional*, defaults to 1): Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is necessary to ensure exact reproducibility of the pretraining results. Please refer to [this issue](https://github.com/pytorch/pytorch/issues/76232). tie_word_embeddings (`bool`, *optional*, defaults to `False`): Whether to tie weight embeddings rope_theta (`float`, *optional*, defaults to 10000.0): The base period of the RoPE embeddings. rope_scaling (`Dict`, *optional*): Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update `max_position_embeddings` to the expected new maximum. attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): Whether to use a bias in the query, key, value and output projection layers during self-attention. attention_dropout (`float`, *optional*, defaults to 0.0): The dropout ratio for the attention probabilities. use_linear_ce (`bool`, *optional*, defaults to `True`): Whether to use cut-cross-entropy (linear CE) for causal LM loss computation when labels are provided. linear_ce_impl (`str`, *optional*, defaults to `"cce_exact"`): Implementation name passed to cut-cross-entropy when `use_linear_ce=True`. ```python >>> from transformers import RizeModel, RizeConfig >>> # Initializing a Rize-V3 style configuration >>> configuration = RizeConfig() >>> # Accessing the model configuration >>> configuration = model.config ```""" model_type = "rize" keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, vocab_size=129280, hidden_size=7168, intermediate_size=18432, moe_intermediate_size = 2048, num_hidden_layers=61, num_nextn_predict_layers=1, num_attention_heads=128, num_key_value_heads=128, n_shared_experts = 1, n_routed_experts = 256, ep_size = 1, routed_scaling_factor = 2.5, kv_lora_rank = 512, q_lora_rank = 1536, qk_rope_head_dim = 64, v_head_dim = 128, qk_nope_head_dim = 128, topk_method = 'noaux_tc', n_group = 8, topk_group = 4, num_experts_per_tok = 8, moe_layer_freq = 1, first_k_dense_replace = 3, norm_topk_prob = True, scoring_func = 'sigmoid', aux_loss_alpha = 0.001, seq_aux = True, auxfree_bias_lr = 0.0, hidden_act="silu", max_position_embeddings=4096, initializer_range=0.02, rms_norm_eps=1e-6, use_cache=True, pad_token_id=None, bos_token_id=0, eos_token_id=1, pretraining_tp=1, tie_word_embeddings=False, rope_theta=10000.0, rope_scaling=None, attention_bias=False, attention_dropout=0.0, use_linear_ce=True, linear_ce_impl="cce_exact", attention_arch="auto", linear_attn_config=None, block_diag_causal_mask=False, reset_position_ids_per_sample=False, moe_router_active_only=True, prompt_loss_weight=0.0, freeze_router_on_sft=False, global_lbl_enabled=False, global_lbl_sync_across_ranks=False, global_lbl_buffer_across_ga=False, **kwargs, ): self.vocab_size = vocab_size self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.moe_intermediate_size = moe_intermediate_size self.num_hidden_layers = num_hidden_layers self.num_nextn_predict_layers = num_nextn_predict_layers self.num_attention_heads = num_attention_heads self.n_shared_experts = n_shared_experts self.n_routed_experts = n_routed_experts self.ep_size = ep_size self.routed_scaling_factor = routed_scaling_factor self.kv_lora_rank = kv_lora_rank self.q_lora_rank = q_lora_rank self.qk_rope_head_dim = qk_rope_head_dim self.v_head_dim = v_head_dim self.qk_nope_head_dim = qk_nope_head_dim self.topk_method = topk_method self.n_group = n_group self.topk_group = topk_group self.num_experts_per_tok = num_experts_per_tok self.moe_layer_freq = moe_layer_freq self.first_k_dense_replace = first_k_dense_replace self.norm_topk_prob = norm_topk_prob self.scoring_func = scoring_func self.aux_loss_alpha = aux_loss_alpha self.seq_aux = seq_aux self.auxfree_bias_lr = auxfree_bias_lr # for backward compatibility if num_key_value_heads is None: num_key_value_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.hidden_act = hidden_act self.initializer_range = initializer_range self.rms_norm_eps = rms_norm_eps self.pretraining_tp = pretraining_tp self.use_cache = use_cache self.rope_theta = rope_theta self.rope_scaling = rope_scaling self.attention_bias = attention_bias self.attention_dropout = attention_dropout self.use_linear_ce = use_linear_ce self.linear_ce_impl = linear_ce_impl self.block_diag_causal_mask = bool(block_diag_causal_mask) self.reset_position_ids_per_sample = bool(reset_position_ids_per_sample) self.moe_router_active_only = bool(moe_router_active_only) self.prompt_loss_weight = float(prompt_loss_weight) self.freeze_router_on_sft = bool(freeze_router_on_sft) self.global_lbl_enabled = bool(global_lbl_enabled) self.global_lbl_sync_across_ranks = bool(global_lbl_sync_across_ranks) self.global_lbl_buffer_across_ga = bool(global_lbl_buffer_across_ga) # ---- Kimi-Linear / hybrid attention knobs (optional) ---- # attention_arch: 'auto' (default; infer from linear_attn_config), 'standard', or 'kimi_linear' self.attention_arch = attention_arch # linear_attn_config: dict or None self.linear_attn_config = linear_attn_config 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, ) @property def is_linear_attn(self) -> bool: """Whether this config enables the Kimi-Linear style hybrid attention. Notes: - This is intentionally a *property* (not a stored boolean) so that it works for both new configs (attention_arch) and older checkpoints that might only carry linear_attn_config. """ arch = getattr(self, "attention_arch", None) la_cfg = getattr(self, "linear_attn_config", None) if arch is None: return la_cfg is not None # allow direct boolean override (e.g. programmatic configs) if isinstance(arch, bool): return bool(arch) and la_cfg is not None arch_s = str(arch).lower().strip() if arch_s in ("auto", "infer"): return la_cfg is not None if arch_s in ("standard", "full", "mla", "default", "none", ""): return False if arch_s in ("kimi_linear", "kda", "linear_attn", "linear", "hybrid"): return la_cfg is not None # Unknown value: be conservative (disable) return False @is_linear_attn.setter def is_linear_attn(self, value: bool) -> None: """Back-compat setter. Allows `config.is_linear_attn = True/False`.""" if bool(value): # Only flip attention_arch if it is currently unset or standard-ish. cur = getattr(self, "attention_arch", "standard") if cur is None or str(cur).lower().strip() in ("standard", "full", "default", "none", ""): self.attention_arch = "kimi_linear" else: self.attention_arch = "standard" def is_kda_layer(self, layer_idx: int) -> bool: """Return True if `layer_idx` (0-indexed) should use KDA (linear attention). Layer index convention: - Internally we expect `layer_idx` to be 0-indexed (as used by `range(num_hidden_layers)`). - In config.linear_attn_config, `kda_layers` / `full_attn_layers` may be either: * 0-indexed (0..num_hidden_layers-1), OR * 1-indexed (1..num_hidden_layers) like the official Kimi-Linear configs. We auto-detect the convention. """ if not self.is_linear_attn: return False cfg = getattr(self, "linear_attn_config", None) if not isinstance(cfg, dict): return False if layer_idx is None: return False try: layer_idx = int(layer_idx) except Exception: return False if layer_idx < 0 or layer_idx >= int(getattr(self, "num_hidden_layers", 0) or 0): return False kda_layers = cfg.get("kda_layers", None) full_layers = cfg.get("full_attn_layers", None) # Normalize lists (best-effort) kda = [int(x) for x in (kda_layers or [])] if isinstance(kda_layers, (list, tuple)) else [] full = [int(x) for x in (full_layers or [])] if isinstance(full_layers, (list, tuple)) else [] # Auto-detect index base for the config lists all_idx = kda + full one_indexed = False if all_idx: n = int(getattr(self, "num_hidden_layers", 0) or 0) # Strong signals: if 0 in all_idx: one_indexed = False elif n in all_idx: one_indexed = True else: # Heuristic: if everything is within [1, n], treat as 1-indexed mn, mx = min(all_idx), max(all_idx) if mn >= 1 and mx <= n: one_indexed = True query_idx = layer_idx + (1 if one_indexed else 0) if kda: return query_idx in set(kda) if full: return query_idx not in set(full) # If neither list is provided, default to "no KDA" return False