from transformers.configuration_utils import PretrainedConfig class MobileLLMP1TextConfig(PretrainedConfig): r""" This is the configuration class to store the configuration of a [`MobileLLMP1TextModel`]. It is used to instantiate a MobileLLMP1 text 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 MobileLLMP1 1B model. 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 202048): Vocabulary size of the Llama4 text model. Defines the maximum number of different tokens that can be represented by the `inputs_ids` passed when calling [`Llama4TextModel`]. hidden_size (`int`, *optional*, defaults to 5120): Dimensionality of the embeddings and hidden states. intermediate_size (`int`, *optional*, defaults to 8192): Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder. intermediate_size_mlp (`int`, *optional*, defaults to 16384): TODO num_hidden_layers (`int`, *optional*, defaults to 48): Number of hidden layers in the Transformer encoder. num_attention_heads (`int`, *optional*, defaults to 40): Number of attention heads for each attention layer in the Transformer encoder. num_key_value_heads (`int`, *optional*, defaults to 8): This is the number of key_value heads that should be used to implement Grouped Query Attention. If not specified, will default to `num_attention_heads`. head_dim (`int`, *optional*, defaults to 128): TODO hidden_act (`str` or `Callable`, *optional*, defaults to `"silu"`): The non-linear activation function (function or string) in the encoder and pooler. max_position_embeddings (`int`, *optional*, defaults to 131072): 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. sliding_window (`int`, *optional*, defaults to 512): In MobileLLMP1, every 4 out of 5 layers use sliding window attention. This is the size of the sliding window. rms_norm_eps (`float`, *optional*, defaults to 1e-05): 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. pad_token_id (`int`, *optional*, defaults to 128004): The id of the padding token. bos_token_id (`int`, *optional*, defaults to 1): The id of the beginning of sentence token. eos_token_id (`int`, *optional*, defaults to 2): The id of the end of sentence token. tie_word_embeddings (`bool`, *optional*, defaults to `False`): Whether to tie weight embeddings rope_theta (`float`, *optional*, defaults to `500000.0`): The base period of the RoPE embeddings. attention_dropout (`int`, *optional*, defaults to 0.0): TODO num_experts_per_tok (`int`, *optional*, defaults to 1): TODO num_local_experts (`int`, *optional*, defaults to 16): TODO moe_layers (`int`, *optional*): TODO interleave_moe_layer_step (`int`, *optional*, defaults to 1): TODO use_qk_norm (`int`, *optional*, defaults to `True`): TODO output_router_logits (`int`, *optional*, defaults to `False`): TODO router_aux_loss_coef (`int`, *optional*, defaults to 0.001): TODO router_jitter_noise (`int`, *optional*, defaults to 0.0): TODO rope_scaling (`Dict`, *optional*): Dictionary containing the scaling configuration for the RoPE embeddings. 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. Expected contents: `rope_type` (`str`): The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope', 'llama3'], with 'default' being the original RoPE implementation. `factor` (`float`, *optional*): Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In most scaling types, a `factor` of x will enable the model to handle sequences of length x * original maximum pre-trained length. `original_max_position_embeddings` (`int`, *optional*): Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during pretraining. `attention_factor` (`float`, *optional*): Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention computation. If unspecified, it defaults to value recommended by the implementation, using the `factor` field to infer the suggested value. `beta_fast` (`float`, *optional*): Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear ramp function. If unspecified, it defaults to 32. `beta_slow` (`float`, *optional*): Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear ramp function. If unspecified, it defaults to 1. `short_factor` (`list[float]`, *optional*): Only used with 'longrope'. The scaling factor to be applied to short contexts (< `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden size divided by the number of attention heads divided by 2 `long_factor` (`list[float]`, *optional*): Only used with 'longrope'. The scaling factor to be applied to long contexts (< `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden size divided by the number of attention heads divided by 2 `low_freq_factor` (`float`, *optional*): Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE `high_freq_factor` (`float`, *optional*): Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE no_rope_layers (`list[int]`, *optional*): List with at least the same length as the number of layers in the model. A `1` at an index position indicates that the corresponding layer will use RoPE, while a `0` indicates that it's a NoPE layer. no_rope_layer_interval (`int`, *optional*, defaults to 4): If `no_rope_layers` is `None`, it will be created using a NoPE layer every `no_rope_layer_interval` layers. attention_chunk_size (`int`, *optional*, defaults to 8192): layer_types (`list`, *optional*): Attention pattern for each layer. attn_temperature_tuning (`bool`, *optional*, defaults to `True`): Whether to dynamically scale the attention temperature for each query token based on sequence length. Recommended for long sequences (e.g., >32k tokens) to maintain stable output results. floor_scale (`int`, *optional*, defaults to 8192): TODO attn_scale (`int`, *optional*, defaults to 0.1): TODO Example: """ model_type = "llama4_text" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise", "layers.*.feed_forward.shared_expert.gate_proj": "local_colwise", "layers.*.feed_forward.shared_expert.up_proj": "local_colwise", "layers.*.feed_forward.shared_expert.down_proj": "local_rowwise", "layers.*.feed_forward.experts.gate_up_proj": "local_packed_rowwise", # row because not linear "layers.*.feed_forward.experts.down_proj": "local_colwise", # col because not linear "layers.*.feed_forward.experts": "local", "layers.*.feed_forward.gate_proj": "local_colwise", "layers.*.feed_forward.up_proj": "local_colwise", "layers.*.feed_forward.down_proj": "local_rowwise", "layers.*.feed_forward": "gather", } base_model_ep_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise", "layers.*.feed_forward.experts.gate_up_proj": "grouped_gemm", # row because not linear "layers.*.feed_forward.experts.down_proj": "grouped_gemm", # col because not linear "layers.*.feed_forward.experts": "gather", # all reduce "layers.*.feed_forward.gate_proj": "local_colwise", "layers.*.feed_forward.up_proj": "local_colwise", "layers.*.feed_forward.down_proj": "local_rowwise", "layers.*.feed_forward.router": "ep_router", } def __init__( self, vocab_size=202048, hidden_size=1280, intermediate_size=6144, intermediate_size_mlp=6144, num_hidden_layers=30, num_attention_heads=20, num_key_value_heads=4, head_dim=64, hidden_act="silu", max_position_embeddings=131072, initializer_range=0.02, rms_norm_eps=1e-5, use_cache=True, pad_token_id=None, sliding_window=512, bos_token_id=1, eos_token_id=2, tie_word_embeddings=True, rope_theta=500000, attention_dropout=0.0, num_experts_per_tok=1, num_local_experts=16, moe_layers=None, interleave_moe_layer_step=1, use_qk_norm=False, output_router_logits=False, router_aux_loss_coef=0.001, router_jitter_noise=0.0, rope_scaling=None, no_rope_layers=None, no_rope_layer_interval=4, attention_chunk_size=8192, layer_types=None, attn_temperature_tuning=True, floor_scale=8192, attn_scale=0.1, **kwargs, ): 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, ) self.attn_temperature_tuning = attn_temperature_tuning self.attn_scale = attn_scale self.floor_scale = floor_scale self.vocab_size = vocab_size self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.intermediate_size_mlp = intermediate_size_mlp self.num_hidden_layers = num_hidden_layers self.sliding_window = sliding_window self.num_attention_heads = num_attention_heads self.rope_scaling = rope_scaling self.attention_bias = False # 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.use_cache = use_cache self.rope_theta = rope_theta self.attention_dropout = attention_dropout self.head_dim = ( head_dim if head_dim is not None else self.hidden_size // self.num_attention_heads ) self.use_qk_norm = use_qk_norm self.num_experts_per_tok = num_experts_per_tok self.num_local_experts = num_local_experts self.output_router_logits = output_router_logits self.router_aux_loss_coef = router_aux_loss_coef self.router_jitter_noise = router_jitter_noise self.layer_types = layer_types # Backwards compatibility if no_rope_layers == []: no_rope_layers = None default_no_rope_layers = [ int((layer_idx + 1) % no_rope_layer_interval != 0) for layer_idx in range(self.num_hidden_layers) ] self.no_rope_layers = ( no_rope_layers if no_rope_layers else default_no_rope_layers ) # If no pattern set, use our default pattern if self.layer_types is None: self.layer_types = [ "sliding_attention" if bool((i) % 4) else "full_attention" for i in range(self.num_hidden_layers) ] + [ "full_attention" ] # Last layer is always full attention self.interleave_moe_layer_step = interleave_moe_layer_step self.moe_layers = ( moe_layers if moe_layers is not None else list( range( interleave_moe_layer_step - 1, num_hidden_layers, interleave_moe_layer_step, ) ) )