| """Configuration for the complete MemT-1.0-Flash model (MIT).""" |
|
|
| import math |
| from copy import deepcopy |
|
|
| from transformers.models.deepseek_v4.configuration_deepseek_v4 import DeepseekV4Config |
|
|
|
|
| _DECLARED_SEMANTICS = { |
| "memory_update_mode": "token", |
| "memory_use_short_conv": False, |
| "memory_head_mode": "mha", |
| "memory_read_scale": "k_inv_sqrt", |
| "tail_policy": "commit", |
| "output_gate_type": "silu", |
| "output_gate_softcap_enabled": True, |
| } |
|
|
|
|
| class MemTDeepseekV4Config(DeepseekV4Config): |
| model_type = "memt_deepseek_v4" |
|
|
| def __init__(self, memt=None, **kwargs): |
| super().__init__(**kwargs) |
| supplied = deepcopy(memt or {}) |
| defaults = { |
| **_DECLARED_SEMANTICS, |
| "layers": [i for i, kind in enumerate(self.layer_types) |
| if i >= 3 and kind == "compressed_sparse_attention"], |
| "projection": "grouped_kv_low_rank", |
| "num_memory_heads": self.num_attention_heads, |
| "num_kv_heads": self.num_attention_heads, |
| "head_dim": self.head_dim, |
| "q_rank": self.q_lora_rank, |
| "kv_groups": 8, |
| "kv_rank": self.head_dim, |
| "gate_rank": self.head_dim, |
| "norm_eps": self.rms_norm_eps, |
| "gate_latent_norm": False, |
| "gdn_version": "stable_gdn2", |
| "stable_gdn2_sqrt_erase_gate": True, |
| "output_gate_bias": True, |
| "output_gate_softcap_limit": 2.0, |
| "chunk_tokens": 2048, |
| } |
| unknown = set(supplied) - set(defaults) |
| if unknown: |
| raise ValueError(f"Unknown MemT fields: {sorted(unknown)}") |
| self.memt = defaults | supplied |
| m = self.memt |
| for name, expected in _DECLARED_SEMANTICS.items(): |
| if type(m[name]) is not type(expected) or m[name] != expected: |
| raise ValueError(f"Unsupported memt.{name}: {m[name]!r}") |
| if (m["projection"] != "grouped_kv_low_rank" |
| or m["gdn_version"] != "stable_gdn2" or m["stable_gdn2_sqrt_erase_gate"] is not True |
| or m["gate_latent_norm"]): |
| raise ValueError("This release implements the trained grouped-KV stable sqrt-b model") |
| if type(m["output_gate_bias"]) is not bool: |
| raise ValueError("memt.output_gate_bias must be bool") |
| limit = m["output_gate_softcap_limit"] |
| if isinstance(limit, bool) or not isinstance(limit, (int, float)) or not math.isfinite(limit) or limit <= 0: |
| raise ValueError("memt.output_gate_softcap_limit must be positive and finite") |
| for name in ("num_memory_heads", "num_kv_heads", "head_dim", "q_rank", "kv_groups", "kv_rank", "gate_rank", "chunk_tokens"): |
| if type(m[name]) is not int or m[name] <= 0: |
| raise ValueError(f"memt.{name} must be a positive integer") |
| if m["num_memory_heads"] % m["kv_groups"]: |
| raise ValueError("MemT heads must be divisible by the number of KV groups") |
| if m["num_kv_heads"] != m["num_memory_heads"]: |
| raise ValueError("MemT KV heads must match memory heads after grouped projection") |
| if m["num_memory_heads"] != self.num_attention_heads or m["head_dim"] != self.head_dim: |
| raise ValueError("MemT output geometry must match the host attention") |
| if len(set(m["layers"])) != len(m["layers"]) or any( |
| not 0 <= i < self.num_hidden_layers or self.layer_types[i] != "compressed_sparse_attention" |
| for i in m["layers"] |
| ): |
| raise ValueError("MemT layers must be distinct CSA layers") |
|
|
| @classmethod |
| def from_dict(cls, config_dict, **kwargs): |
| result = super().from_dict(deepcopy(config_dict), **kwargs) |
| config = result[0] if isinstance(result, tuple) else result |
| |
| config.memt = cls(**config.to_dict()).memt |
| return result |
|
|
|
|
| MemTDeepseekV4Config.register_for_auto_class() |
|
|
