| |
| """Configuration for ConvGPT-v2 hybrid 1D/2D causal language model.""" |
|
|
| from __future__ import annotations |
|
|
| from typing import Optional, Sequence |
|
|
| from transformers import PretrainedConfig |
|
|
|
|
| class ConvGPTV2Config(PretrainedConfig): |
| """ |
| Transformers-compatible config for a hybrid causal ConvGPT-v2 language model. |
| |
| The model mixes a causal 1D branch with a causal 2D branch over a configurable |
| square token grid. 2D packing is configurable and includes row-major, snake, |
| Morton/Z-order, and Hilbert mappings. Tiny retrieval/router blocks can be |
| inserted every N layers for cheap long-range content routing. |
| """ |
|
|
| model_type = "convgpt_v2" |
| keys_to_ignore_at_inference = ["past_key_values"] |
|
|
| def __init__( |
| self, |
| vocab_size: int = 32768, |
| hidden_size: int = 512, |
| intermediate_size: Optional[int] = None, |
| num_hidden_layers: int = 12, |
| max_position_embeddings: Optional[int] = None, |
| grid_size: int = 128, |
| packing: str = "hilbert", |
| pack_order: str = "sequence_to_curve", |
| conv1d_kernel_size: int = 5, |
| conv1d_dilations: Optional[Sequence[int]] = None, |
| conv2d_kernel_size: int = 3, |
| conv2d_dilations: Optional[Sequence[int]] = None, |
| conv_expand: int = 2, |
| conv2d_backend: str = "chunked_gather", |
| conv2d_chunk_size: int = 1024, |
| use_1d_branch: bool = True, |
| use_2d_branch: bool = True, |
| use_2d_depthwise: bool = True, |
| two_d_every: int = 1, |
| two_d_start_layer: int = 0, |
| conv1d_residual_gate_init: float = 0.0, |
| conv2d_residual_gate_init: float = -4.0, |
| branch_dropout: float = 0.0, |
| position_embedding_type: str = "learned", |
| rope_theta: float = 10000.0, |
| router_rope_fraction: float = 1.0, |
| use_row_col_embeddings: bool = True, |
| fusion: str = "gated", |
| retrieval_every: int = 4, |
| retrieval_num_slots: int = 64, |
| retrieval_top_k: int = 4, |
| retrieval_num_heads: int = 4, |
| retrieval_dropout: Optional[float] = None, |
| router_type: str = "topk_memory", |
| chunk_memory_size: int = 64, |
| chunk_memory_top_k: int = 4, |
| chunk_memory_token_top_k: int = 0, |
| chunk_memory_use_triton: bool = True, |
| chunk_memory_gate_init: float = -4.0, |
| chunk_memory_include_current_chunk: bool = False, |
| hidden_act: str = "silu", |
| rms_norm_eps: float = 1e-6, |
| dropout: float = 0.0, |
| initializer_range: float = 0.02, |
| tie_word_embeddings: bool = True, |
| use_cache: bool = False, |
| pad_token_id: int = 1, |
| bos_token_id: int = 0, |
| eos_token_id: int = 2, |
| **kwargs, |
| ): |
| if grid_size <= 0: |
| raise ValueError("grid_size must be positive") |
| if hidden_size <= 0: |
| raise ValueError("hidden_size must be positive") |
| if num_hidden_layers <= 0: |
| raise ValueError("num_hidden_layers must be positive") |
| if conv1d_kernel_size <= 0 or conv1d_kernel_size % 2 == 0: |
| raise ValueError("conv1d_kernel_size must be a positive odd integer") |
| if conv2d_kernel_size <= 0 or conv2d_kernel_size % 2 == 0: |
| raise ValueError("conv2d_kernel_size must be a positive odd integer") |
|
|
| packing = packing.lower().replace("-", "_") |
| aliases = {"z_order": "morton", "zorder": "morton", "hilbert_curve": "hilbert"} |
| packing = aliases.get(packing, packing) |
| if packing not in {"row_major", "snake", "morton", "hilbert"}: |
| raise ValueError( |
| "packing must be one of: row_major, snake, morton/z_order, hilbert" |
| ) |
|
|
| pack_order = pack_order.lower() |
| if pack_order not in {"sequence_to_curve", "curve_to_sequence"}: |
| raise ValueError("pack_order must be sequence_to_curve or curve_to_sequence") |
|
|
| fusion = fusion.lower() |
| if fusion not in {"gated", "sum", "concat"}: |
| raise ValueError("fusion must be gated, sum, or concat") |
|
|
| position_embedding_type = position_embedding_type.lower().replace("-", "_") |
| if position_embedding_type not in {"learned", "nope", "rope_nope"}: |
| raise ValueError("position_embedding_type must be learned, nope, or rope_nope") |
| if rope_theta <= 0: |
| raise ValueError("rope_theta must be positive") |
| if not 0.0 <= router_rope_fraction <= 1.0: |
| raise ValueError("router_rope_fraction must be in [0, 1]") |
|
|
| router_type = router_type.lower() |
| if router_type not in {"topk_memory", "chunk_memory", "chunk_token_memory", "none"}: |
| raise ValueError("router_type must be topk_memory, chunk_memory, chunk_token_memory, or none") |
| if chunk_memory_size <= 0: |
| raise ValueError("chunk_memory_size must be positive") |
| if chunk_memory_top_k <= 0: |
| raise ValueError("chunk_memory_top_k must be positive") |
| if chunk_memory_token_top_k < 0: |
| raise ValueError("chunk_memory_token_top_k must be non-negative") |
|
|
| conv2d_backend = conv2d_backend.lower() |
| if conv2d_backend not in {"unfold", "chunked_gather", "masked_conv2d", "triton_gather"}: |
| raise ValueError("conv2d_backend must be unfold, chunked_gather, masked_conv2d, or triton_gather") |
| if conv2d_backend == "masked_conv2d" and packing != "row_major": |
| raise ValueError("conv2d_backend=masked_conv2d is exact-causal only for row_major packing") |
| if conv2d_chunk_size <= 0: |
| raise ValueError("conv2d_chunk_size must be positive") |
| if two_d_every <= 0: |
| raise ValueError("two_d_every must be positive") |
| if two_d_start_layer < 0: |
| raise ValueError("two_d_start_layer must be non-negative") |
|
|
| max_tokens = grid_size * grid_size |
| if max_position_embeddings is None: |
| max_position_embeddings = max_tokens |
| if max_position_embeddings > max_tokens: |
| raise ValueError( |
| f"max_position_embeddings ({max_position_embeddings}) cannot exceed grid_size^2 ({max_tokens})" |
| ) |
|
|
| self.vocab_size = vocab_size |
| self.hidden_size = hidden_size |
| self.intermediate_size = intermediate_size if intermediate_size is not None else hidden_size * 4 |
| self.num_hidden_layers = num_hidden_layers |
| self.max_position_embeddings = max_position_embeddings |
| self.grid_size = grid_size |
| self.max_grid_tokens = max_tokens |
| self.packing = packing |
| self.pack_order = pack_order |
| self.conv1d_kernel_size = conv1d_kernel_size |
| self.conv1d_dilations = list(conv1d_dilations) if conv1d_dilations is not None else None |
| self.conv2d_kernel_size = conv2d_kernel_size |
| self.conv2d_dilations = list(conv2d_dilations) if conv2d_dilations is not None else None |
| self.conv_expand = conv_expand |
| self.conv2d_backend = conv2d_backend |
| self.conv2d_chunk_size = conv2d_chunk_size |
| self.use_1d_branch = use_1d_branch |
| self.use_2d_branch = use_2d_branch |
| self.use_2d_depthwise = use_2d_depthwise |
| self.two_d_every = two_d_every |
| self.two_d_start_layer = two_d_start_layer |
| self.conv1d_residual_gate_init = conv1d_residual_gate_init |
| self.conv2d_residual_gate_init = conv2d_residual_gate_init |
| self.branch_dropout = branch_dropout |
| self.position_embedding_type = position_embedding_type |
| self.rope_theta = rope_theta |
| self.router_rope_fraction = router_rope_fraction |
| self.use_row_col_embeddings = use_row_col_embeddings |
| self.fusion = fusion |
| self.retrieval_every = retrieval_every |
| self.retrieval_num_slots = retrieval_num_slots |
| self.retrieval_top_k = retrieval_top_k |
| self.retrieval_num_heads = retrieval_num_heads |
| self.retrieval_dropout = dropout if retrieval_dropout is None else retrieval_dropout |
| self.router_type = router_type |
| self.chunk_memory_size = chunk_memory_size |
| self.chunk_memory_top_k = chunk_memory_top_k |
| self.chunk_memory_token_top_k = chunk_memory_token_top_k |
| self.chunk_memory_gate_init = chunk_memory_gate_init |
| self.chunk_memory_include_current_chunk = chunk_memory_include_current_chunk |
| self.hidden_act = hidden_act |
| self.rms_norm_eps = rms_norm_eps |
| self.dropout = dropout |
| self.initializer_range = initializer_range |
| self.tie_word_embeddings = tie_word_embeddings |
| self.use_cache = use_cache |
| self.is_encoder_decoder = False |
|
|
| 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, |
| ) |
|
|
|
|
| __all__ = ["ConvGPTV2Config"] |
|
|
