configuration_llama4.py 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479
  1. # coding=utf-8
  2. # Copyright 2025 The LLAMA4 and HuggingFace Inc. team. All rights reserved.
  3. #
  4. #
  5. # Licensed under the Apache License, Version 2.0 (the "License");
  6. # you may not use this file except in compliance with the License.
  7. # You may obtain a copy of the License at
  8. #
  9. # http://www.apache.org/licenses/LICENSE-2.0
  10. #
  11. # Unless required by applicable law or agreed to in writing, software
  12. # distributed under the License is distributed on an "AS IS" BASIS,
  13. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14. # See the License for the specific language governing permissions and
  15. # limitations under the License.
  16. import warnings
  17. from ...configuration_utils import PretrainedConfig, layer_type_validation
  18. from ...utils import logging
  19. logger = logging.get_logger(__name__)
  20. class Llama4VisionConfig(PretrainedConfig):
  21. r"""
  22. This is the configuration class to store the configuration of a [`Llama4VisionModel`]. It is used to instantiate a
  23. Llama4 vision model according to the specified arguments, defining the model architecture. Instantiating a configuration
  24. with the defaults will yield a similar configuration to that of the Llama4 109B.
  25. e.g. [meta-llama/Llama-4-Scout-17B-16E](https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E)
  26. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
  27. documentation from [`PretrainedConfig`] for more information.
  28. Args:
  29. hidden_size (`int`, *optional*, defaults to 768):
  30. Dimensionality of the encoder layers and the pooler layer.
  31. hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
  32. The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
  33. `"relu"`, `"selu"` and `"gelu_new"` `"quick_gelu"` are supported.
  34. num_hidden_layers (`int`, *optional*, defaults to 34):
  35. Number of hidden layers in the Transformer encoder.
  36. num_attention_heads (`int`, *optional*, defaults to 16):
  37. Number of attention heads for each attention layer in the Transformer encoder.
  38. num_channels (`int`, *optional*, defaults to 3):
  39. Number of channels in the input image.
  40. intermediate_size (`int`, *optional*, defaults to 5632):
  41. Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
  42. vision_output_dim (`int`, *optional*, defaults to 7680):
  43. Dimensionality of the vision model output. Includes output of transformer
  44. encoder with intermediate layers and global transformer encoder.
  45. image_size (`int`, *optional*, defaults to 448):
  46. The size (resolution) of each image *tile*.
  47. patch_size (`int`, *optional*, defaults to 14):
  48. The size (resolution) of each patch.
  49. norm_eps (`float`, *optional*, defaults to 1e-05):
  50. The epsilon used by the layer normalization layers.
  51. vision_feature_select_strategy (`int`, *optional*, defaults to `"default"`): TODO
  52. initializer_range (`float`, *optional*, defaults to 0.02):
  53. The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
  54. pixel_shuffle_ratio (`int`, *optional*, defaults to 0.5): TODO
  55. projector_input_dim (`int`, *optional*, defaults to 4096): TODO
  56. projector_output_dim (`int`, *optional*, defaults to 4096): TODO
  57. multi_modal_projector_bias (`int`, *optional*, defaults to `False`): TODO
  58. projector_dropout (`int`, *optional*, defaults to 0.0): TODO
  59. attention_dropout (`int`, *optional*, defaults to 0.0): TODO
  60. rope_theta (`int`, *optional*, defaults to 10000): TODO
  61. """
  62. base_model_tp_plan = {
  63. "model.layers.*.self_attn.q_proj": "colwise",
  64. "model.layers.*.self_attn.k_proj": "colwise",
  65. "model.layers.*.self_attn.v_proj": "colwise",
  66. "model.layers.*.self_attn.o_proj": "rowwise",
  67. "vision_adapter.mlp.fc1": "colwise",
  68. "vision_adapter.mlp.fc2": "rowwise",
  69. "patch_embedding.linear": "colwise_rep",
  70. }
  71. model_type = "llama4_vision_model"
  72. base_config_key = "vision_config"
  73. def __init__(
  74. self,
  75. hidden_size: int = 768,
  76. hidden_act: str = "gelu",
  77. num_hidden_layers: int = 34,
  78. num_attention_heads: int = 16,
  79. num_channels: int = 3,
  80. intermediate_size: int = 5632,
  81. vision_output_dim: int = 7680,
  82. image_size: int = 448,
  83. patch_size: int = 14,
  84. norm_eps: float = 1e-5,
  85. vision_feature_select_strategy="default",
  86. initializer_range: float = 0.02,
  87. pixel_shuffle_ratio=0.5,
  88. projector_input_dim=4096,
  89. projector_output_dim=4096,
  90. multi_modal_projector_bias=False,
  91. projector_dropout=0.0,
  92. attention_dropout=0.0,
  93. rope_theta=10000,
  94. **kwargs,
  95. ):
  96. self.hidden_size = hidden_size
  97. self.hidden_act = hidden_act
  98. self.num_hidden_layers = num_hidden_layers
  99. self.num_channels = num_channels
  100. self.intermediate_size = intermediate_size
  101. self.image_size = image_size
  102. self.vision_output_dim = vision_output_dim
  103. self.patch_size = patch_size
  104. self.norm_eps = norm_eps
  105. self.num_attention_heads = num_attention_heads
  106. self.initializer_range = initializer_range
  107. self.pixel_shuffle_ratio = pixel_shuffle_ratio
  108. self.projector_input_dim = projector_input_dim
  109. self.projector_output_dim = projector_output_dim
  110. self.multi_modal_projector_bias = multi_modal_projector_bias
  111. self.projector_dropout = projector_dropout
  112. self.attention_dropout = attention_dropout
  113. self.vision_feature_select_strategy = vision_feature_select_strategy
  114. self.rope_theta = rope_theta
  115. self._vision_feature_layer = kwargs.get("vision_feature_layer", -1)
  116. @property
  117. def vision_feature_layer(self):
  118. warnings.warn(
  119. "The `vision_feature_layer` attribute is deprecated and will be removed in v4.58.0.",
  120. FutureWarning,
  121. )
  122. return self._vision_feature_layer
  123. @vision_feature_layer.setter
  124. def vision_feature_layer(self, value):
  125. self._vision_feature_layer = value
  126. super().__init__(**kwargs)
  127. class Llama4TextConfig(PretrainedConfig):
  128. r"""
  129. This is the configuration class to store the configuration of a [`Llama4TextModel`]. It is used to instantiate a
  130. Llama4 text model according to the specified arguments, defining the model architecture. Instantiating a configuration
  131. with the defaults will yield a similar configuration to that of the Llama4 109B.
  132. e.g. [meta-llama/Llama-4-Scout-17B-16E](https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E)
  133. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
  134. documentation from [`PretrainedConfig`] for more information.
  135. Args:
  136. vocab_size (`int`, *optional*, defaults to 202048):
  137. Vocabulary size of the Llama4 text model. Defines the maximum number of different tokens that can be represented
  138. by the `inputs_ids` passed when calling [`Llama4TextModel`].
  139. hidden_size (`int`, *optional*, defaults to 5120):
  140. Dimensionality of the embeddings and hidden states.
  141. intermediate_size (`int`, *optional*, defaults to 8192):
  142. Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
  143. intermediate_size_mlp (`int`, *optional*, defaults to 16384): TODO
  144. num_hidden_layers (`int`, *optional*, defaults to 48):
  145. Number of hidden layers in the Transformer encoder.
  146. num_attention_heads (`int`, *optional*, defaults to 40):
  147. Number of attention heads for each attention layer in the Transformer encoder.
  148. num_key_value_heads (`int`, *optional*, defaults to 8):
  149. This is the number of key_value heads that should be used to implement Grouped Query Attention. If not
  150. specified, will default to `num_attention_heads`.
  151. head_dim (`int`, *optional*, defaults to 128): TODO
  152. hidden_act (`str` or `Callable`, *optional*, defaults to `"silu"`):
  153. The non-linear activation function (function or string) in the encoder and pooler.
  154. max_position_embeddings (`int`, *optional*, defaults to 131072):
  155. The maximum sequence length that this model might ever be used with.
  156. initializer_range (`float`, *optional*, defaults to 0.02):
  157. The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
  158. rms_norm_eps (`float`, *optional*, defaults to 1e-05):
  159. The epsilon used by the rms normalization layers.
  160. use_cache (`bool`, *optional*, defaults to `True`):
  161. Whether or not the model should return the last key/values attentions.
  162. pad_token_id (`int`, *optional*, defaults to 128004):
  163. The id of the padding token.
  164. bos_token_id (`int`, *optional*, defaults to 1):
  165. The id of the beginning of sentence token.
  166. eos_token_id (`int`, *optional*, defaults to 2):
  167. The id of the end of sentence token.
  168. tie_word_embeddings (`bool`, *optional*, defaults to `False`):
  169. Whether to tie weight embeddings
  170. rope_theta (`float`, *optional*, defaults to `500000.0`):
  171. The base period of the RoPE embeddings.
  172. attention_dropout (`int`, *optional*, defaults to 0.0): TODO
  173. num_experts_per_tok (`int`, *optional*, defaults to 1): TODO
  174. num_local_experts (`int`, *optional*, defaults to 16): TODO
  175. moe_layers (`int`, *optional*): TODO
  176. interleave_moe_layer_step (`int`, *optional*, defaults to 1): TODO
  177. use_qk_norm (`int`, *optional*, defaults to `True`): TODO
  178. output_router_logits (`int`, *optional*, defaults to `False`): TODO
  179. router_aux_loss_coef (`int`, *optional*, defaults to 0.001): TODO
  180. router_jitter_noise (`int`, *optional*, defaults to 0.0): TODO
  181. rope_scaling (`Dict`, *optional*):
  182. Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
  183. and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
  184. accordingly.
  185. Expected contents:
  186. `rope_type` (`str`):
  187. The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
  188. 'llama3'], with 'default' being the original RoPE implementation.
  189. `factor` (`float`, *optional*):
  190. Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
  191. most scaling types, a `factor` of x will enable the model to handle sequences of length x *
  192. original maximum pre-trained length.
  193. `original_max_position_embeddings` (`int`, *optional*):
  194. Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
  195. pretraining.
  196. `attention_factor` (`float`, *optional*):
  197. Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
  198. computation. If unspecified, it defaults to value recommended by the implementation, using the
  199. `factor` field to infer the suggested value.
  200. `beta_fast` (`float`, *optional*):
  201. Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
  202. ramp function. If unspecified, it defaults to 32.
  203. `beta_slow` (`float`, *optional*):
  204. Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
  205. ramp function. If unspecified, it defaults to 1.
  206. `short_factor` (`list[float]`, *optional*):
  207. Only used with 'longrope'. The scaling factor to be applied to short contexts (<
  208. `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
  209. size divided by the number of attention heads divided by 2
  210. `long_factor` (`list[float]`, *optional*):
  211. Only used with 'longrope'. The scaling factor to be applied to long contexts (<
  212. `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
  213. size divided by the number of attention heads divided by 2
  214. `low_freq_factor` (`float`, *optional*):
  215. Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
  216. `high_freq_factor` (`float`, *optional*):
  217. Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
  218. <TODO>
  219. <TODO>
  220. no_rope_layers (`list[int]`, *optional*):
  221. List with at least the same length as the number of layers in the model.
  222. A `1` at an index position indicates that the corresponding layer will use RoPE,
  223. while a `0` indicates that it's a NoPE layer.
  224. no_rope_layer_interval (`int`, *optional*, defaults to 4):
  225. If `no_rope_layers` is `None`, it will be created using a NoPE layer every
  226. `no_rope_layer_interval` layers.
  227. attention_chunk_size (`int`, *optional*, defaults to 8192):
  228. <TODO>
  229. layer_types (`list`, *optional*):
  230. Attention pattern for each layer.
  231. attn_temperature_tuning (`bool`, *optional*, defaults to `True`):
  232. Whether to dynamically scale the attention temperature for each query token based on sequence length.
  233. Recommended for long sequences (e.g., >32k tokens) to maintain stable output results.
  234. floor_scale (`int`, *optional*, defaults to 8192): TODO
  235. attn_scale (`int`, *optional*, defaults to 0.1): TODO
  236. Example:
  237. """
  238. model_type = "llama4_text"
  239. keys_to_ignore_at_inference = ["past_key_values"]
  240. base_model_tp_plan = {
  241. "layers.*.self_attn.q_proj": "colwise",
  242. "layers.*.self_attn.k_proj": "colwise",
  243. "layers.*.self_attn.v_proj": "colwise",
  244. "layers.*.self_attn.o_proj": "rowwise",
  245. "layers.*.feed_forward.shared_expert.gate_proj": "local_colwise",
  246. "layers.*.feed_forward.shared_expert.up_proj": "local_colwise",
  247. "layers.*.feed_forward.shared_expert.down_proj": "local_rowwise",
  248. "layers.*.feed_forward.experts.gate_up_proj": "local_packed_rowwise", # row because not linear
  249. "layers.*.feed_forward.experts.down_proj": "local_colwise", # col because not linear
  250. "layers.*.feed_forward.experts": "local",
  251. "layers.*.feed_forward.gate_proj": "local_colwise",
  252. "layers.*.feed_forward.up_proj": "local_colwise",
  253. "layers.*.feed_forward.down_proj": "local_rowwise",
  254. "layers.*.feed_forward": "gather",
  255. }
  256. base_model_ep_plan = {
  257. "layers.*.self_attn.q_proj": "colwise",
  258. "layers.*.self_attn.k_proj": "colwise",
  259. "layers.*.self_attn.v_proj": "colwise",
  260. "layers.*.self_attn.o_proj": "rowwise",
  261. "layers.*.feed_forward.experts.gate_up_proj": "grouped_gemm", # row because not linear
  262. "layers.*.feed_forward.experts.down_proj": "grouped_gemm", # col because not linear
  263. "layers.*.feed_forward.experts": "gather", # all reduce
  264. "layers.*.feed_forward.gate_proj": "local_colwise",
  265. "layers.*.feed_forward.up_proj": "local_colwise",
  266. "layers.*.feed_forward.down_proj": "local_rowwise",
  267. "layers.*.feed_forward.router": "ep_router",
  268. }
  269. def __init__(
  270. self,
  271. vocab_size=202048,
  272. hidden_size=5120,
  273. intermediate_size=8192,
  274. intermediate_size_mlp=16384,
  275. num_hidden_layers=48,
  276. num_attention_heads=40,
  277. num_key_value_heads=8,
  278. head_dim=128,
  279. hidden_act="silu",
  280. max_position_embeddings=4096 * 32,
  281. initializer_range=0.02,
  282. rms_norm_eps=1e-5,
  283. use_cache=True,
  284. pad_token_id=None,
  285. bos_token_id=1,
  286. eos_token_id=2,
  287. tie_word_embeddings=False,
  288. rope_theta=500000,
  289. attention_dropout=0.0,
  290. num_experts_per_tok=1,
  291. num_local_experts=16,
  292. moe_layers=None,
  293. interleave_moe_layer_step=1,
  294. use_qk_norm=True,
  295. output_router_logits=False,
  296. router_aux_loss_coef=0.001,
  297. router_jitter_noise=0.0,
  298. rope_scaling=None,
  299. no_rope_layers=None,
  300. no_rope_layer_interval=4,
  301. attention_chunk_size=8192,
  302. layer_types=None,
  303. attn_temperature_tuning=True,
  304. floor_scale=8192,
  305. attn_scale=0.1,
  306. **kwargs,
  307. ):
  308. super().__init__(
  309. pad_token_id=pad_token_id,
  310. bos_token_id=bos_token_id,
  311. eos_token_id=eos_token_id,
  312. tie_word_embeddings=tie_word_embeddings,
  313. **kwargs,
  314. )
  315. self.attn_temperature_tuning = attn_temperature_tuning
  316. self.attn_scale = attn_scale
  317. self.floor_scale = floor_scale
  318. self.vocab_size = vocab_size
  319. self.max_position_embeddings = max_position_embeddings
  320. self.hidden_size = hidden_size
  321. self.intermediate_size = intermediate_size
  322. self.intermediate_size_mlp = intermediate_size_mlp
  323. self.num_hidden_layers = num_hidden_layers
  324. self.num_attention_heads = num_attention_heads
  325. self.rope_scaling = rope_scaling
  326. self.attention_bias = False
  327. # for backward compatibility
  328. if num_key_value_heads is None:
  329. num_key_value_heads = num_attention_heads
  330. self.num_key_value_heads = num_key_value_heads
  331. self.hidden_act = hidden_act
  332. self.initializer_range = initializer_range
  333. self.rms_norm_eps = rms_norm_eps
  334. self.use_cache = use_cache
  335. self.rope_theta = rope_theta
  336. self.attention_dropout = attention_dropout
  337. self.head_dim = head_dim if head_dim is not None else self.hidden_size // self.num_attention_heads
  338. self.use_qk_norm = use_qk_norm
  339. self.num_experts_per_tok = num_experts_per_tok
  340. self.num_local_experts = num_local_experts
  341. self.output_router_logits = output_router_logits
  342. self.router_aux_loss_coef = router_aux_loss_coef
  343. self.router_jitter_noise = router_jitter_noise
  344. # Backwards compatibility
  345. if no_rope_layers == []:
  346. no_rope_layers = None
  347. default_no_rope_layers = [
  348. int((layer_idx + 1) % no_rope_layer_interval != 0) for layer_idx in range(self.num_hidden_layers)
  349. ]
  350. self.no_rope_layers = no_rope_layers if no_rope_layers else default_no_rope_layers
  351. self.interleave_moe_layer_step = interleave_moe_layer_step
  352. self.moe_layers = (
  353. moe_layers
  354. if moe_layers is not None
  355. else list(range(interleave_moe_layer_step - 1, num_hidden_layers, interleave_moe_layer_step))
  356. )
  357. self.attention_chunk_size = attention_chunk_size
  358. self.layer_types = layer_types
  359. if layer_types is None:
  360. self.layer_types = [
  361. "chunked_attention" if no_rope else "full_attention" for no_rope in self.no_rope_layers
  362. ]
  363. layer_type_validation(self.layer_types, self.num_hidden_layers)
  364. class Llama4Config(PretrainedConfig):
  365. r"""
  366. This is the configuration class to store the configuration of a [`Llama4Model`]. It is used to instantiate an
  367. Llama4 model according to the specified arguments, defining the model architecture. Instantiating a configuration
  368. with the defaults will yield a similar configuration to that of the Llama4 109B.
  369. e.g. [meta-llama/Llama-4-Scout-17B-16E](https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E)
  370. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
  371. documentation from [`PretrainedConfig`] for more information.
  372. Args:
  373. vision_config (`Llama4VisionConfig`, *optional*):
  374. The Llama4 Vision config.
  375. text_config (`Llama4TextConfig`, *optional*):
  376. The Llama4 Text config.
  377. boi_token_index (`int`, *optional*, defaults to 200080):
  378. The begin-of-image token index to wrap the image prompt.
  379. eoi_token_index (`int`, *optional*, defaults to 200081):
  380. The end-of-image token index to wrap the image prompt.
  381. image_token_index (`int`, *optional*, defaults to 200092):
  382. The image token index to encode the image prompt.
  383. tie_word_embeddings (`bool`, *optional*, defaults to `False`):
  384. Whether the model's input and output word embeddings should be tied.
  385. ```python
  386. >>> from transformers import Llama4Model, Llama4Config
  387. >>> # Initializing a Llama4 7B style configuration
  388. >>> configuration = Llama4Config()
  389. >>> # Initializing a model from the Llama4 7B style configuration
  390. >>> model = Llama4Model(configuration)
  391. >>> # Accessing the model configuration
  392. >>> configuration = model.config
  393. ```"""
  394. model_type = "llama4"
  395. attribute_map = {
  396. "image_token_id": "image_token_index",
  397. "boi_token_id": "boi_token_index",
  398. "eoi_token_id": "eoi_token_index",
  399. }
  400. sub_configs = {"text_config": Llama4TextConfig, "vision_config": Llama4VisionConfig}
  401. base_model_tp_plan = {
  402. "multi_modal_projector.linear_1": "colwise_rep",
  403. }
  404. def __init__(
  405. self,
  406. vision_config=None,
  407. text_config=None,
  408. boi_token_index=200080,
  409. eoi_token_index=200081,
  410. image_token_index=200092,
  411. tie_word_embeddings=False,
  412. **kwargs,
  413. ):
  414. if vision_config is None:
  415. self.vision_config = Llama4VisionConfig()
  416. logger.info("vision_config is None, using default llama4 vision config")
  417. elif isinstance(vision_config, dict):
  418. self.vision_config = Llama4VisionConfig(**vision_config)
  419. elif isinstance(vision_config, Llama4VisionConfig):
  420. self.vision_config = vision_config
  421. self.boi_token_index = boi_token_index
  422. self.eoi_token_index = eoi_token_index
  423. self.image_token_index = image_token_index
  424. if text_config is None:
  425. self.text_config = Llama4TextConfig()
  426. logger.info("text_config is None, using default llama4 text config")
  427. elif isinstance(text_config, dict):
  428. self.text_config = Llama4TextConfig(**text_config)
  429. elif isinstance(text_config, Llama4TextConfig):
  430. self.text_config = text_config
  431. super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
  432. __all__ = ["Llama4Config", "Llama4TextConfig", "Llama4VisionConfig"]