configuration_helium1_casa.py

from typing import Any, Literal

from transformers.configuration_utils import PretrainedConfig
from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLVisionConfig


class Helium1CASAConfig(PretrainedConfig):
    r"""
    Helium1 Config augmented with CASA options


    Args:
        vocab_size (`int`, *optional*, defaults to 32000):
            Vocabulary size of the LLaMA model. Defines the number of different tokens that can be represented by the
            `inputs_ids` passed when calling [`Helium1Model`]
        hidden_size (`int`, *optional*, defaults to 4096):
            Dimension of the hidden representations.
        intermediate_size (`int`, *optional*, defaults to 11008):
            Dimension of the MLP representations.
        num_hidden_layers (`int`, *optional*, defaults to 32):
            Number of hidden layers in the Transformer decoder.
        num_attention_heads (`int`, *optional*, defaults to 32):
            Number of attention heads for each attention layer in the Transformer decoder.
        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. Llama 1 supports up to 2048 tokens,
            Llama 2 up to 4096, CodeLlama up to 16384.
        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/main/perf_train_gpu_many#tensor-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. 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
        attention_bias (`bool`, *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.
        mlp_bias (`bool`, *optional*, defaults to `False`):
            Whether to use a bias in up_proj, down_proj and gate_proj layers in the MLP layers.
        head_dim (`int`, *optional*):
            The attention head dimension. If None, it will default to hidden_size // num_attention_heads

    """

    model_type = "helium1_casa"
    keys_to_ignore_at_inference = ["past_key_values"]
    # Default tensor parallel plan for base model `Helium1Model`
    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.*.mlp.gate_proj": "colwise",
        "layers.*.mlp.up_proj": "colwise",
        "layers.*.mlp.down_proj": "rowwise",
    }
    base_model_pp_plan = {  # pyright: ignore[reportAssignmentType]
        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
        "norm": (["hidden_states"], ["hidden_states"]),
    }

    def __init__(
        self,
        vocab_size: int = 32000,
        hidden_size: int = 4096,
        intermediate_size: int = 11008,
        num_hidden_layers: int = 32,
        num_attention_heads: int = 32,
        num_key_value_heads: None | int = None,
        head_dim: None | int = None,
        hidden_act: str = "silu",
        attention_dropout: float = 0.0,
        max_position_embeddings: int = 2048,
        initializer_range: float = 0.02,
        rms_norm_eps: float = 1e-6,
        use_cache: bool = True,
        tie_word_embeddings: bool = False,
        rope_theta: float = 10000.0,
        pad_token_id: int = 3,
        eos_token_id: int = 2,
        bos_token_id: int = 1,
        pretraining_tp: int = 1,
        rope_scaling: None | dict = None,
        attention_bias: bool = False,
        mlp_bias: bool = False,
        # Our fusion mechanisms
        # Common to all fusion mechanisms
        xa_layers: None | tuple = None,
        xa_order: Literal["ca_first", "parallel", "instead"] = "ca_first",
        xa_norm_on_images: bool = False,
        xa_update_image_embeds: bool = False,
        mask_squash_blockwise: bool = False,
        # CASA
        casa_attention: bool = False,
        casa_delta_w: bool = False,
        casa_windows: Literal["batch", "squashed", "images", "turn_based"] = "batch",
        casa_use_asymetric_qkv: bool = True,
        xa_custom_norm: bool = False,
        # Qwen2.5-VL vision config
        vision_config: dict[str, Any] | None = None,
        **kwargs: Any,
    ):
        from transformers.modeling_rope_utils import rope_config_validation

        self.vocab_size = vocab_size
        self.max_position_embeddings = max_position_embeddings
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads

        # 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.head_dim = (
            head_dim if head_dim is not None else self.hidden_size // self.num_attention_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.mlp_bias = mlp_bias
        # Validate the correctness of rotary position embeddings parameters
        # BC: if there is a 'type' field, copy it it to 'rope_type'.
        if self.rope_scaling is not None and "type" in self.rope_scaling:
            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
        rope_config_validation(self)

        self.head_dim = self.hidden_size // self.num_attention_heads
        self.xa_layers = xa_layers
        self.xa_order: Literal["ca_first", "parallel", "instead"] = xa_order
        self.xa_norm_on_images = xa_norm_on_images
        self.xa_update_image_embeds = xa_update_image_embeds
        self.mask_squash_blockwise = mask_squash_blockwise
        # CASA config
        self.casa_attention = casa_attention
        self.casa_delta_w = casa_delta_w
        self.casa_windows: Literal["batch", "squashed", "images", "turn_based"] = casa_windows
        self.casa_use_asymetric_qkv = casa_use_asymetric_qkv
        self.xa_custom_norm = xa_custom_norm

        if vision_config is None:
            vision_config = dict()
        self.vision_config = Qwen2_5_VLVisionConfig(**vision_config)
        self.vision_config.temporal_patch_size = 1
        self.vision_config.image_mean = [0.48145466, 0.4578275, 0.40821073]
        self.vision_config.image_std = [0.26862954, 0.26130258, 0.27577711]
        self.vision_config.out_dim = 2048

        self.pre_image_tokens = []
        self.post_image_tokens = []

        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,
        )


if __name__ == "__main__":
    import argparse
    from pathlib import Path

    import rich
    import yaml
    from transformers.models.auto.configuration_auto import AutoConfig

    parser = argparse.ArgumentParser()
    parser.add_argument("--out_dir", type=str, default="./saved_config/")
    parser.add_argument(
        "--ckpt_path",
        type=str,
        default="/lustre/scwpod02/client/kyutai/juliette/experiments/finext_casa_896_xtxt_up_b20_64gpu/fdf76e6774",
    )
    args = parser.parse_args()
    path = Path(args.ckpt_path) / "kyuteye_config.yml"

    helium_config = AutoConfig.from_pretrained("kyutai/helium-1-2b")
    vision_config = AutoConfig.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct").vision_config

    # 3) Create YOUR config by merging both
    config = Helium1CASAConfig(
        **helium_config.to_dict(),  # all helium parameters
        vision_config=vision_config.to_dict(),  # override or add vision_config
    )

    with open(path) as stream:
        kconfig = yaml.safe_load(stream)

    # print keys that are in kconfig and in config
    for key in set(kconfig.keys()).intersection(set(config.to_dict().keys())):
        rich.print(f"Overwriting [bold green]{key:>50s}[/]: [bold red]{kconfig[key]}")
        setattr(config, key, kconfig[key])
    # TODO: handle casa_own_norm -> xa_custom_norm
    print("Configuration successfully loaded.")
    # Save config to json
    config.save_pretrained(args.out_dir)
    print(f"Configuration saved to {args.out_dir}/config.json")