| """HuggingFace configuration for Logos.""" |
|
|
| from __future__ import annotations |
|
|
| from types import SimpleNamespace |
| from typing import Any, Dict |
|
|
| from transformers import PretrainedConfig |
|
|
|
|
| _NATIVE_DEFAULTS: Dict[str, Any] = { |
| "vocab_size": 32000, |
| "d_model": 512, |
| "max_seq_len": 2048, |
| "num_layers": 8, |
| "num_heads": 8, |
| "norm_eps": 1e-6, |
| "d_ff": 1364, |
| "use_moe": True, |
| "num_shared_experts": 2, |
| "num_sparse_experts": 64, |
| "top_k": 6, |
| "expert_d_ff": 256, |
| "bias_update_rate": 0.01, |
| "capacity_factor": 2.0, |
| "router_logit_noise_std": 0.1, |
| "router_logit_noise_decay_steps": 2000, |
| "router_init_std": 0.002, |
| "router_bias_error_clip": 1.0, |
| "router_bias_clip": 1.0, |
| "moe_aux_loss_weight": 1e-3, |
| "moe_aux_loss_decay_steps": 2000, |
| "lm_head_chunk_size": 0, |
| "moe_diversity_factor": 0.0, |
| "rope_base": 10000.0, |
| "qk_norm": True, |
| "partial_rope_dim": None, |
| "attention_sink": True, |
| "block_residual_isolate_softmax": False, |
| "head_dim": 64, |
| "conv_size": 4, |
| "chunk_size": 64, |
| "A_init_range": (1, 16), |
| "expand": 2, |
| "swa_window": 256, |
| "swa_every": 4, |
| "swa_offset": 3, |
| "csa_compression": 4, |
| "csa_top_k": 1024, |
| "csa_indexer_heads": 4, |
| "csa_indexer_dim": 32, |
| "csa_indexer_loss_weight": 1.0, |
| "hca_compression": 128, |
| "compressed_query_dim": None, |
| "compressed_head_dim": None, |
| "compressed_rope": False, |
| "indexer_rope": False, |
| "attn_pattern": None, |
| "num_entry_layers": 2, |
| "num_body_layers": 4, |
| "num_exit_layers": 2, |
| "num_loops": 4, |
| "entry_attn_pattern": None, |
| "body_attn_pattern": None, |
| "exit_attn_pattern": None, |
| "entry_top_k": None, |
| "exit_top_k": None, |
| "gradient_checkpointing": False, |
| "ckpt_granularity": "per-block", |
| } |
|
|
|
|
| class LogosConfig(PretrainedConfig): |
| model_type = "logos" |
| keys_to_ignore_at_inference = ["past_key_values"] |
|
|
| def __init__(self, **kwargs: Any): |
| native_values: Dict[str, Any] = {} |
| for name, default in _NATIVE_DEFAULTS.items(): |
| native_values[name] = kwargs.pop(name, default) |
|
|
| tokenizer_encoding = kwargs.pop("tokenizer_encoding", "cl100k_base") |
| kwargs.setdefault("tie_word_embeddings", True) |
| super().__init__(**kwargs) |
|
|
| for name, value in native_values.items(): |
| setattr(self, name, value) |
| self.tokenizer_encoding = tokenizer_encoding |
| self.architectures = ["LogosForCausalLM"] |
| self.auto_map = { |
| "AutoConfig": "configuration_logos.LogosConfig", |
| "AutoModelForCausalLM": "modeling_logos.LogosForCausalLM", |
| "AutoTokenizer": "tokenization_logos.LogosTokenizer", |
| } |
|
|
| def to_native_config(self): |
| data = {name: getattr(self, name) for name in _NATIVE_DEFAULTS} |
| data["num_layers"] = ( |
| int(data["num_entry_layers"]) |
| + int(data["num_body_layers"]) |
| + int(data["num_exit_layers"]) |
| ) |
| return SimpleNamespace(**data) |
|
|
|
|
| __all__ = ["LogosConfig"] |
|
|
