configuration_mistral4.py 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149
  1. # Copyright 2026 Mistral AI and The HuggingFace Inc. team. All rights reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. """Mistral4 model configuration"""
  15. from huggingface_hub.dataclasses import strict
  16. from ...configuration_utils import PreTrainedConfig
  17. from ...modeling_rope_utils import RopeParameters
  18. from ...utils import auto_docstring
  19. @auto_docstring(checkpoint="mistralai/Mistral-Small-4-119B-2603")
  20. @strict
  21. class Mistral4Config(PreTrainedConfig):
  22. r"""
  23. n_group (`int`, *optional*, defaults to 1):
  24. Number of groups for routed experts.
  25. first_k_dense_replace (`int`, *optional*, defaults to 0):
  26. Number of dense layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head).
  27. \--k dense layers--/
  28. rope_interleave (`bool`, *optional*, defaults to `True`):
  29. Whether to interleave the rotary position embeddings.
  30. Example:
  31. ```python
  32. >>> from transformers import Mistral4Model, Mistral4Config
  33. >>> # Initializing a Mistral4 style configuration
  34. >>> configuration = Mistral4Config()
  35. >>> # Accessing the model configuration
  36. >>> configuration = model.config
  37. ```"""
  38. model_type = "mistral4"
  39. keys_to_ignore_at_inference = ["past_key_values"]
  40. base_model_tp_plan = {
  41. "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
  42. "layers.*.mlp.experts.down_proj": "rowwise",
  43. "layers.*.mlp.experts": "moe_tp_experts",
  44. "layers.*.mlp.shared_experts.gate_proj": "colwise",
  45. "layers.*.mlp.shared_experts.up_proj": "colwise",
  46. "layers.*.mlp.shared_experts.down_proj": "rowwise",
  47. "layers.*.mlp.gate_proj": "colwise",
  48. "layers.*.mlp.up_proj": "colwise",
  49. "layers.*.mlp.down_proj": "rowwise",
  50. }
  51. base_model_pp_plan = {
  52. "embed_tokens": (["input_ids"], ["inputs_embeds"]),
  53. "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
  54. "norm": (["hidden_states"], ["hidden_states"]),
  55. }
  56. attribute_map = {
  57. "num_local_experts": "n_routed_experts",
  58. }
  59. vocab_size: int = 131072
  60. hidden_size: int = 4096
  61. intermediate_size: int = 12288
  62. moe_intermediate_size: int = 2048
  63. num_hidden_layers: int = 36
  64. num_attention_heads: int = 32
  65. num_key_value_heads: int | None = 32
  66. n_shared_experts: int = 1
  67. n_routed_experts: int = 128
  68. routed_scaling_factor: float = 1.0
  69. kv_lora_rank: int = 256
  70. q_lora_rank: int | None = 1024
  71. qk_rope_head_dim: int = 64
  72. v_head_dim: int | None = 128
  73. qk_nope_head_dim: int = 64
  74. n_group: int | None = 1
  75. topk_group: int | None = 1
  76. num_experts_per_tok: int | None = 4
  77. first_k_dense_replace: int | None = 0
  78. norm_topk_prob: bool | None = True
  79. hidden_act: str = "silu"
  80. max_position_embeddings: int = 1048576
  81. initializer_range: float = 0.02
  82. rms_norm_eps: float = 1e-6
  83. use_cache: bool = True
  84. pad_token_id: int | None = 11
  85. bos_token_id: int | None = 1
  86. eos_token_id: int | list[int] | None = 2
  87. pretraining_tp: int | None = 1
  88. tie_word_embeddings: bool = False
  89. rope_parameters: RopeParameters | dict | None = None
  90. rope_interleave: bool | None = True
  91. attention_bias: bool = False
  92. attention_dropout: float | int | None = 0.0
  93. def __post_init__(self, **kwargs):
  94. if self.rope_parameters is None:
  95. self.rope_parameters = {
  96. "type": "yarn",
  97. "rope_theta": 10000.0,
  98. "factor": 128.0,
  99. "original_max_position_embeddings": 8192,
  100. "max_position_embeddings": self.max_position_embeddings,
  101. "beta_fast": 32.0,
  102. "beta_slow": 1.0,
  103. "mscale_all_dim": 1.0,
  104. "mscale": 1.0,
  105. "llama_4_scaling_beta": 0.1,
  106. "partial_rotary_factor": self.qk_rope_head_dim / (self.qk_nope_head_dim + self.qk_rope_head_dim),
  107. }
  108. if self.num_key_value_heads is None:
  109. self.num_key_value_heads = self.num_attention_heads
  110. self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
  111. self.head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
  112. self.rope_parameters.setdefault("partial_rotary_factor", self.qk_rope_head_dim / self.head_dim)
  113. super().__post_init__(
  114. ignore_keys_at_rope_validation={"llama_4_scaling_beta", "max_position_embeddings"}, **kwargs
  115. )
  116. def convert_rope_params_to_dict(self, ignore_keys_at_rope_validation: set | None = None, **kwargs):
  117. rope_scaling = kwargs.pop("rope_scaling", None)
  118. self.rope_parameters = rope_scaling or self.rope_parameters
  119. self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else {}
  120. # Standardize and validate the correctness of rotary position embeddings parameters
  121. self.rope_parameters.setdefault("rope_theta", kwargs.pop("rope_theta", self.default_theta))
  122. self.standardize_rope_params()
  123. if ignore_keys_at_rope_validation is not None:
  124. self.ignore_keys_at_rope_validation = self.ignore_keys_at_rope_validation | ignore_keys_at_rope_validation
  125. self.validate_rope()
  126. # Convert to float because RoPE fn expect a float. Models on the hub were saved as int
  127. for key in ["beta_fast", "beta_slow", "factor"]:
  128. if key in self.rope_parameters:
  129. self.rope_parameters[key] = float(self.rope_parameters[key])
  130. return kwargs
  131. __all__ = ["Mistral4Config"]