Repository: huggingface/transformers
File: src/transformers/models/llama4/configuration_llama4.py (Python)

Code before:
            layers alternate. Increasing this step reduces the total number of experts, making the model more dense-like,
            which can improve training stability but reduces the total parameter capacity.
        use_qk_norm (`bool`, *optional*, defaults to `True`):
            Whether to normalize the Query and Key matrices in the attention layer. This stabilizes training and preventing
            attention logits from growing too large and causing numerical instability.
        output_router_logits (`bool`, *optional*, defaults to `False`):
            Whether or not to return the router logits of all MoE layers.
        router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
            The aux loss factor for the total loss. If it's too high the model may prioritize balanced expert and if it's
            too low the model may suffer from expert collapse.
        router_jitter_noise (`float`, *optional*, defaults to 0.0):
            The amount of noise to add to the router logits. This encourages exploration, preventing the router
            from prematurely converging to a single expert (expert collapse). For pure inference, it's often set to 0 or very low.
        rope_parameters (`RopeParameters`, *optional*):
            Dictionary containing the configuration parameters for the RoPE embeddings. The dictionary should contain
            a value for `rope_theta` and optionally parameters used for scaling in case you want to use RoPE
            with longer `max_position_embeddings`.
        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):
            Chunk size for the attention computation.
        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):
            Scaling factor for the floor operation in the attention mechanism. A higher scale (like 8192) results in
            "smoother" or "diffuse" distributions.
        attn_scale (`float`, *optional*, defaults to 0.1):
            Scaling factor for the attention scores. A lower scale (like 0.1) creates "sharper" attention distributions,
            while a higher scale (like 1.0) results in "smoother" or "diffuse" distributions.

    Example:
    """

    model_type = "llama4_text"
    keys_to_ignore_at_inference = ["past_key_values"]
    default_theta = 500000.0
    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",

Review comment (suggestion):
```suggestion
            Chunk size for the attention computation. Smaller value enforces more local attention and lowers memory.
```

Code after:
        rope_parameters (`RopeParameters`, *optional*):
            Dictionary containing the configuration parameters for the RoPE embeddings. The dictionary should contain
            a value for `rope_theta` and optionally parameters used for scaling in case you want to use RoPE
            with longer `max_position_embeddings`.
        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):
            Chunk size for the attention computation. Smaller value enforces more local attention and lowers memory.
        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):
            Base scale (in tokens) for attention temperature tuning. Larger value delays scaling to longer positions.
        attn_scale (`float`, *optional*, defaults to 0.1):
            Strength of attention temperature tuning. Larger value increases scaling at long positions.

    Example:
    """

    model_type = "llama4_text"
    keys_to_ignore_at_inference = ["past_key_values"]
    default_theta = 500000.0
    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": "colwise",
        "layers.*.feed_forward.shared_expert.up_proj": "colwise",
        "layers.*.feed_forward.shared_expert.down_proj": "rowwise",
        "layers.*.feed_forward.experts.gate_up_proj": "packed_rowwise",  # row because not linear
        "layers.*.feed_forward.experts.down_proj": "colwise",  # col because not linear
        "layers.*.feed_forward.gate_proj": "colwise",
        "layers.*.feed_forward.up_proj": "colwise",
        "layers.*.feed_forward.down_proj": "rowwise",
    }
    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.gate_proj": "colwise",