| from __future__ import annotations |
|
|
| from transformers.configuration_utils import PretrainedConfig |
| from transformers.utils import logging |
|
|
| logger = logging.get_logger(__name__) |
|
|
| class AfmoeConfig(PretrainedConfig): |
| """ |
| n_group (`int`, *optional*, defaults to 1): |
| Number of groups for routed experts. |
| topk_group (`int`, *optional*, defaults to 1): |
| Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups). |
| """ |
| model_type = "afmoe" |
| base_model_pp_plan = { |
| "embed_tokens": (["input_ids"], ["inputs_embeds"]), |
| "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), |
| "norm": (["hidden_states"], ["hidden_states"]), |
| } |
|
|
| def __init__( |
| self, |
| num_hidden_layers: int = 32, |
| vocab_size: int = 200192, |
| hidden_size: int = 2048, |
| intermediate_size: int = 6144, |
| moe_intermediate_size=1408, |
| num_dense_layers=1, |
| num_attention_heads=16, |
| num_key_value_heads=None, |
| head_dim=128, |
| hidden_act="silu", |
| max_position_embeddings=16384, |
| initializer_range=0.02, |
| rms_norm_eps=1e-5, |
| use_cache=True, |
| tie_word_embeddings=False, |
| rope_theta=10000.0, |
| rope_scaling=None, |
| num_experts=64, |
| num_experts_per_tok=6, |
| num_shared_experts=2, |
| num_expert_groups=1, |
| num_limited_groups=1, |
| score_func="sigmoid", |
| route_norm=True, |
| route_scale=1.0, |
| global_attn_every_n_layers=4, |
| sliding_window=1024, |
| mup_enabled=False, |
| layer_types=None, |
| attention_dropout: float = 0.0, |
| n_group: int = 1, |
| topk_group: int = 1, |
| **kwargs, |
| ): |
| 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_dense_layers = num_dense_layers |
| self.num_attention_heads = num_attention_heads |
| self.head_dim = head_dim |
| 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.rope_scaling = rope_scaling |
| self.rope_parameters = dict(rope_scaling) if rope_scaling is not None else None |
| |
| |
| |
| self.moe_intermediate_size = moe_intermediate_size |
| self.num_experts_per_tok = num_experts_per_tok |
| self.n_group = n_group |
| self.topk_group = topk_group |
| self.num_experts = num_experts |
| self.num_shared_experts = num_shared_experts |
| self.num_expert_groups = num_expert_groups |
| self.num_limited_groups = num_limited_groups |
| self.score_func = score_func |
| self.route_norm = route_norm |
| self.route_scale = route_scale |
|
|
|
|
| |
| self.attention_dropout = attention_dropout |
| self.global_attn_every_n_layers = global_attn_every_n_layers |
| self.sliding_window = sliding_window |
| self.layer_types = layer_types |
| if self.layer_types is None: |
| self.layer_types = [ |
| "sliding_attention" if bool((i + 1) % global_attn_every_n_layers) else "full_attention" for i in range(self.num_hidden_layers) |
| ] |
| self.validate_layer_type() |
|
|
| |
| self.mup_enabled = mup_enabled |
|
|
| if num_key_value_heads is None: |
| num_key_value_heads = num_attention_heads |
|
|
| self.num_key_value_heads = num_key_value_heads |
|
|
|
|
| |
| if self.rope_scaling is not None and "type" in self.rope_scaling: |
| self.rope_scaling["rope_type"] = self.rope_scaling["type"] |
| if self.rope_parameters is not None and "type" in self.rope_parameters: |
| self.rope_parameters["rope_type"] = self.rope_parameters["type"] |
| self.standardize_rope_params() |
| self.validate_rope() |
|
|
| super().__init__( |
| tie_word_embeddings=tie_word_embeddings, |
| **kwargs, |
| ) |
| class TrinityVLMConfig(PretrainedConfig): |
| model_type = "trinity_vlm" |
| is_composition = True |
|
|
| def __init__( |
| self, |
| text_config: dict | None = None, |
| vision_config: dict | None = None, |
| projector_hidden_dim: int | None = None, |
| vision_feature_dim: int | None = None, |
| image_seq_len: int | None = None, |
| enable_grouped_moe: bool = True, |
| output_router_logits: bool = False, |
| image_start_token: str = "<|vision_start|>", |
| image_end_token: str = "<|vision_end|>", |
| image_token: str = "<|image_pad|>", |
| image_start_token_id: int | None = None, |
| image_end_token_id: int | None = None, |
| image_token_id: int | None = None, |
| hidden_size: int | None = None, |
| vocab_size: int | None = None, |
| **kwargs, |
| ) -> None: |
| text_config = dict(text_config or {}) |
| vision_config = dict( |
| vision_config |
| or { |
| "enc_dim": 1152, |
| "enc_patch_size": 14, |
| "enc_n_layers": 27, |
| "enc_ff_dim": 4304, |
| "enc_n_heads": 16, |
| "proj_out_dim": 2048, |
| "crop_size": 378, |
| "in_channels": 3, |
| "max_crops": 12, |
| "overlap_margin": 4, |
| "proj_inner_dim": 8192, |
| "projector_hidden_dim": 2048, |
| } |
| ) |
| if projector_hidden_dim is not None: |
| vision_config["projector_hidden_dim"] = projector_hidden_dim |
|
|
| if hidden_size is None: |
| hidden_size = text_config.get("hidden_size") |
| if vocab_size is None: |
| vocab_size = text_config.get("vocab_size") |
| kwargs.setdefault("bos_token_id", text_config.get("bos_token_id")) |
| kwargs.setdefault("eos_token_id", text_config.get("eos_token_id")) |
| kwargs.setdefault("pad_token_id", text_config.get("pad_token_id")) |
| if vision_feature_dim is None: |
| vision_feature_dim = vision_config.get("proj_out_dim", 2048) |
| if image_seq_len is None: |
| crop_size = int(vision_config.get("crop_size", 378)) |
| patch_size = int(vision_config.get("enc_patch_size", 14)) |
| image_seq_len = (crop_size // patch_size) ** 2 |
|
|
| self.text_config = text_config |
| self.vision_config = vision_config |
| self.vision_feature_dim = vision_feature_dim |
| self.image_seq_len = image_seq_len |
| self.enable_grouped_moe = enable_grouped_moe |
| self.output_router_logits = output_router_logits |
| self.image_start_token = image_start_token |
| self.image_end_token = image_end_token |
| self.image_token = image_token |
| self.image_start_token_id = image_start_token_id |
| self.image_end_token_id = image_end_token_id |
| self.image_token_id = image_token_id |
| self.hidden_size = hidden_size |
| self.vocab_size = vocab_size |
| super().__init__(**kwargs) |
|
|
| @property |
| def projector_hidden_dim(self) -> int: |
| return int(self.vision_config.get("projector_hidden_dim", self.vision_feature_dim)) |
|
|
| def get_text_config(self, decoder: bool = False): |
| del decoder |
| text_config = AfmoeConfig(**self.text_config) |
| text_config.vocab_size = self.vocab_size |
| text_config.bos_token_id = self.bos_token_id |
| text_config.eos_token_id = self.eos_token_id |
| text_config.pad_token_id = self.pad_token_id |
| text_config.packed_experts = True |
| text_config.enable_grouped_moe = self.enable_grouped_moe |
| text_config.output_router_logits = self.output_router_logits |
| return text_config |
|
|
|
|
| __all__ = ["AfmoeConfig", "TrinityVLMConfig"] |
|
|
