| from dataclasses import asdict, dataclass |
| from typing import Any, Dict |
|
|
|
|
| @dataclass |
| class MabaSparseConfig: |
| dim: int = 640 |
| d_model: int = 640 |
| n_layers: int = 16 |
| num_layers: int = 16 |
| n_heads: int = 10 |
| num_heads: int = 10 |
| d_head: int = 64 |
| d_k: int = 64 |
| d_v: int = 64 |
| vocab_size: int = 32000 |
| max_seq_len: int = 2048 |
| d_emb: int = 320 |
| kernel_size: int = 4 |
| conv_kernel_size: int = 4 |
| chunk_size: int = 16 |
| eps: float = 1e-6 |
| inversion_method: str = "adaptive" |
| inversion_threshold: float = 0.5 |
| adaptive_tol: float = 7e-5 |
| d_c: int = 128 |
| d_idx: int = 64 |
| block_size: int = 64 |
| window_size: int = 128 |
| top_k: int = 32 |
| hca_pool_size: int = 64 |
| dist_lambda: float = 0.5 |
| intermediate_size: int = 1728 |
| rms_norm_eps: float = 1e-6 |
| residual_gate_bias: float = 2.0 |
| mtp_depth: int = 2 |
|
|
| def __post_init__(self) -> None: |
| if self.dim != 640 and self.d_model == 640: |
| self.d_model = self.dim |
| elif self.d_model != 640 and self.dim == 640: |
| self.dim = self.d_model |
|
|
| if self.n_heads != 10 and self.num_heads == 10: |
| self.num_heads = self.n_heads |
| elif self.num_heads != 10 and self.n_heads == 10: |
| self.n_heads = self.num_heads |
|
|
| if self.n_layers != 16 and self.num_layers == 16: |
| self.num_layers = self.n_layers |
| elif self.num_layers != 16 and self.n_layers == 16: |
| self.n_layers = self.num_layers |
|
|
| if self.kernel_size != 4 and self.conv_kernel_size == 4: |
| self.conv_kernel_size = self.kernel_size |
| elif self.conv_kernel_size != 4 and self.kernel_size == 4: |
| self.kernel_size = self.conv_kernel_size |
|
|
| if self.d_head == 64 and self.dim != 640: |
| self.d_head = self.dim // self.n_heads |
| self.d_k = self.d_head |
| self.d_v = self.d_head |
|
|
| valid_methods = {"adaptive", "neumann", "exact"} |
| if not isinstance(self.inversion_method, str) or self.inversion_method.lower() not in valid_methods: |
| raise ValueError( |
| f"Invalid inversion_method: '{self.inversion_method}'. " |
| f"Supported options: ['adaptive', 'neumann', 'exact']." |
| ) |
| self.inversion_method = self.inversion_method.lower() |
|
|
| def to_dict(self) -> Dict[str, Any]: |
| return asdict(self) |
|
|
| @classmethod |
| def from_dict(cls, d: Dict[str, Any]) -> "MabaSparseConfig": |
| return cls(**d) |
|
|
