| """LuauCoder configuration (``trust_remote_code=True``).
|
|
|
| Mirrors ``lm.config.ModelConfig`` with transformers' standard names.
|
| ``num_attention_heads`` counts attention output heads: with ``diff_v2`` the model
|
| runs twice as many query heads against ``num_key_value_heads`` K/V heads.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| import math
|
|
|
| from transformers import PretrainedConfig
|
|
|
| __all__ = ["LuauCoderConfig"]
|
|
|
|
|
| class LuauCoderConfig(PretrainedConfig):
|
| model_type = "luau_coder"
|
| keys_to_ignore_at_inference = ["past_key_values"]
|
|
|
| def __init__(
|
| self,
|
| vocab_size: int = 32768,
|
| hidden_size: int = 2560,
|
| intermediate_size: int = 6144,
|
| num_hidden_layers: int = 36,
|
| num_attention_heads: int = 20,
|
| num_key_value_heads: int = 5,
|
| head_dim: int = 128,
|
| rms_norm_eps: float = 1e-6,
|
| final_logit_softcapping: float | None = None,
|
| max_position_embeddings: int = 4096,
|
| rope_theta: float = 500000.0,
|
| rope_scaling: dict | None = None,
|
| lerope_alpha: list[float] | None = None,
|
| nope_layers: list[int] | None = None,
|
| diff_v2: bool = True,
|
| xsa: bool = True,
|
| qk_norm: bool = True,
|
| use_output_gate: bool | str | None = False,
|
| attn_res: bool = False,
|
| attn_res_block_size: int = 6,
|
| attn_res_delta_block: bool = False,
|
| attn_res_partial_source: bool = False,
|
| ccgqa: bool = False,
|
| ccgqa_seq_kernel: int = 2,
|
| ccgqa_head_kernel: int = 2,
|
| ccgqa_qk_mean: bool = True,
|
| ccgqa_value_shift: bool = True,
|
| vha: bool = False,
|
| iha: bool = False,
|
| iha_pseudo_heads: int = 2,
|
| iha_window: int = 0,
|
| iha_global_every: int = 5,
|
| head_gate_rank: int = 0,
|
| initializer_range: float | None = None,
|
| use_cache: bool = True,
|
| tie_word_embeddings: bool = False,
|
|
|
| fused_kernels: bool = True,
|
| compile_elementwise: bool = True,
|
| loss_chunk_size: int = 4096,
|
| bos_token_id: int | None = None,
|
| eos_token_id: int | list[int] | None = 32744,
|
| pad_token_id: int | None = 32745,
|
| **kwargs,
|
| ):
|
| self.vocab_size = vocab_size
|
| self.hidden_size = hidden_size
|
| self.intermediate_size = intermediate_size
|
| self.num_hidden_layers = num_hidden_layers
|
| self.num_attention_heads = num_attention_heads
|
| self.num_key_value_heads = num_key_value_heads
|
| self.head_dim = head_dim
|
| self.rms_norm_eps = rms_norm_eps
|
| self.final_logit_softcapping = final_logit_softcapping
|
| self.max_position_embeddings = max_position_embeddings
|
| self.rope_theta = rope_theta
|
| self.rope_scaling = rope_scaling
|
|
|
| self.lerope_alpha = None if lerope_alpha is None else [float(a) for a in lerope_alpha]
|
| self.nope_layers = sorted(nope_layers or [])
|
| self.diff_v2 = diff_v2
|
| self.xsa = xsa
|
| self.qk_norm = qk_norm
|
|
|
| self.use_output_gate = ("elementwise" if use_output_gate is True
|
| else str(use_output_gate).lower() if use_output_gate else False)
|
| self.attn_res = attn_res
|
| self.attn_res_block_size = attn_res_block_size
|
| self.attn_res_delta_block = attn_res_delta_block
|
| self.attn_res_partial_source = attn_res_partial_source
|
| self.ccgqa = bool(ccgqa)
|
| self.ccgqa_seq_kernel = ccgqa_seq_kernel
|
| self.ccgqa_head_kernel = ccgqa_head_kernel
|
| self.ccgqa_qk_mean = ccgqa_qk_mean
|
| self.ccgqa_value_shift = ccgqa_value_shift
|
| self.vha = bool(vha)
|
| self.iha = bool(iha)
|
| self.iha_pseudo_heads = iha_pseudo_heads
|
| self.iha_window = iha_window
|
| self.iha_global_every = iha_global_every
|
|
|
| self.head_gate_rank = int(head_gate_rank)
|
| if self.head_gate_rank < 0:
|
| raise ValueError("head_gate_rank must be >= 0")
|
| if self.iha and (iha_pseudo_heads < 1 or iha_window < 0 or iha_global_every < 0):
|
| raise ValueError("iha: iha_pseudo_heads >= 1, iha_window >= 0 and iha_global_every >= 0")
|
| self.initializer_range = hidden_size ** -0.5 if initializer_range is None else initializer_range
|
| self.use_cache = use_cache
|
| self.fused_kernels = fused_kernels
|
| self.compile_elementwise = compile_elementwise
|
| self.loss_chunk_size = loss_chunk_size
|
| if num_attention_heads % num_key_value_heads:
|
| raise ValueError(f"num_attention_heads ({num_attention_heads}) must be a multiple of "
|
| f"num_key_value_heads ({num_key_value_heads})")
|
| if head_dim % 2:
|
| raise ValueError("head_dim must be even (RoPE)")
|
| if self.use_output_gate not in (False, "elementwise", "headwise"):
|
| raise ValueError(f"use_output_gate: true / 'elementwise', 'headwise' or false, got {use_output_gate!r}")
|
| if self.ccgqa and not qk_norm:
|
| raise ValueError("ccgqa needs qk_norm")
|
| if self.ccgqa and ccgqa_value_shift and (num_key_value_heads * head_dim) % 2:
|
| raise ValueError("ccgqa_value_shift needs an even num_key_value_heads * head_dim")
|
| if attn_res and math.ceil(num_hidden_layers / attn_res_block_size) > 32:
|
| raise ValueError("attn_res: at most 32 blocks (num_hidden_layers / attn_res_block_size)")
|
| if any(not 0 <= i < num_hidden_layers for i in self.nope_layers):
|
| raise ValueError(f"nope_layers must be in [0, {num_hidden_layers}), got {self.nope_layers}")
|
| if rope_scaling is not None:
|
| rtype = rope_scaling.get("rope_type", rope_scaling.get("type"))
|
| if rtype not in (None, "default", "yarn"):
|
| raise ValueError(f"rope_scaling: only 'yarn' is supported, got {rtype!r}")
|
| super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id,
|
| pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs)
|
|
|
| @property
|
| def output_gate(self) -> str | None:
|
| """None (no gate), "elementwise" or "headwise"."""
|
| return self.use_output_gate or None
|
|
|
| @property
|
| def num_query_heads(self) -> int:
|
| """Query heads the attention kernel runs (2x the output heads with DiffAttn V2)."""
|
| return self.num_attention_heads * (2 if self.diff_v2 else 1)
|
|
|
| @property
|
| def num_q_latent_heads(self) -> int:
|
| """Query heads ``qkv`` produces (VHA: the physical heads every KV group shares)."""
|
| return self.num_query_heads // self.num_key_value_heads if self.vha else self.num_query_heads
|
|
|
| def iha_layer(self, i: int) -> bool:
|
| """Whether layer ``i`` runs IHA (the rest are plain global attention layers)."""
|
| every = self.iha_global_every
|
| return self.iha and not (every and (i + 1) % every == 0)
|
|
|
