""" Configuration class for Q-TensorFormer: Closed-Loop Information-Adaptive Tensor Network LLM. Supports: - Grouped-Query Attention (GQA) with RoPE - Tensor-Train (TT) low-rank linear layers - Closed-loop PID dynamic resource controller - Adaptive KV cache with attention sinks & cold-line quantization - Turnkey Hugging Face AutoConfig registration """ import math from typing import List, Optional, Tuple, Union from transformers.configuration_utils import PretrainedConfig from transformers import AutoConfig def factorize_dim(dim: int, num_cores: int = 3) -> Tuple[int, ...]: """ Heuristically factorize an integer dimension into `num_cores` balanced sub-dimensions. Used for automatic TT core shape inference. """ factors: List[int] = [] d = dim for i in range(2, int(math.isqrt(dim)) + 1): while d % i == 0: factors.append(i) d //= i if d > 1: factors.append(d) if not factors: factors = [dim] # Combine or pad to exactly num_cores while len(factors) < num_cores: # Split the largest factor idx = max(range(len(factors)), key=lambda i: factors[i]) val = factors[idx] sqrt_val = math.isqrt(val) if sqrt_val > 1 and val % sqrt_val == 0: factors[idx] = sqrt_val factors.append(val // sqrt_val) else: factors.append(1) while len(factors) > num_cores: # Multiply the two smallest factors factors.sort() factors[0] = factors[0] * factors[1] factors.pop(1) factors.sort(reverse=True) return tuple(factors) class QTensorFormerConfig(PretrainedConfig): r""" This is the configuration class to store the configuration of a [`QTensorFormerForCausalLM`]. It is used to instantiate a Q-TensorFormer model according to the specified arguments, defining the model architecture, Tensor-Train decomposition parameters, closed-loop PID controller, and adaptive KV cache. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. Args: vocab_size (`int`, *optional*, defaults to 32000): Vocabulary size of the Q-TensorFormer model. Defines the number of different tokens that can be represented by the `inputs_ids` passed when calling [`QTensorFormerForCausalLM`]. hidden_size (`int`, *optional*, defaults to 2048): Dimension of the hidden representations. num_hidden_layers (`int`, *optional*, defaults to 16): Number of hidden layers in the Transformer decoder. num_attention_heads (`int`, *optional*, defaults to 16): Number of attention heads for each attention layer in the Transformer decoder. num_key_value_heads (`int`, *optional*, defaults to 4): This is the number of key_value heads that should be used to implement Grouped Query Attention. If `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. intermediate_size (`int`, *optional*, defaults to 5632): Dimension of the MLP representations. max_position_embeddings (`int`, *optional*, defaults to 4096): The maximum sequence length that this model might ever be used with. rms_norm_eps (`float`, *optional*, defaults to 1e-6): The epsilon used by the rms normalization layers. tt_ranks (`List[int]`, *optional*, defaults to `[1, 16, 16, 1]`): Tensor-Train ranks $[R_0, R_1, \dots, R_d]$ for linear projections. Must start and end with 1. tt_shapes (`List[Tuple[Tuple[int, ...], Tuple[int, ...]]]`, *optional*): Explicit list of input/output factor shapes for TT projections. If None, automatically inferred. enable_dynamic_routing (`bool`, *optional*, defaults to `True`): Whether to dynamically modulate TT contraction rank and KV cache precision on a per-token basis. pid_kp (`float`, *optional*, defaults to 0.1): Proportional gain for the closed-loop dual-subgradient PID resource controller. pid_ki (`float`, *optional*, defaults to 0.01): Integral gain for the closed-loop dual-subgradient PID resource controller. pid_kd (`float`, *optional*, defaults to 0.05): Derivative gain for the closed-loop dual-subgradient PID resource controller. hysteresis_threshold (`float`, *optional*, defaults to 0.15): Hysteresis deadband $\tau$ to prevent high-frequency routing chattering and oscillation. attention_sink_size (`int`, *optional*, defaults to 4): Number of initial tokens to permanently preserve in the KV cache as attention sinks. kv_cache_window_size (`int`, *optional*, defaults to 1024): Sliding window size for local token residency in the adaptive KV cache. quantize_kv_cold (`bool`, *optional*, defaults to `True`): Whether to dynamically quantize cold KV tokens outside the immediate window. kv_quant_bits (`int`, *optional*, defaults to 4): Quantization bit-width for cold KV tokens (4 or 8). target_energy_budget (`float`, *optional*, defaults to 1.0): Normalized energy budget set-point for the PID controller. rope_theta (`float`, *optional*, defaults to 10000.0): The base period of the RoPE embeddings. initializer_range (`float`, *optional*, defaults to 0.02): The standard deviation of the truncated_normal_initializer for initializing all weight matrices. tie_word_embeddings (`bool`, *optional*, defaults to `False`): Whether to tie weight embeddings. """ model_type = "qtensorformer" keys_to_ignore_at_loading = [] def __init__( self, vocab_size: int = 32000, hidden_size: int = 2048, num_hidden_layers: int = 16, num_attention_heads: int = 16, num_key_value_heads: Optional[int] = None, intermediate_size: int = 5632, max_position_embeddings: int = 4096, rms_norm_eps: float = 1e-6, tt_ranks: Optional[List[int]] = None, tt_shapes: Optional[List[Tuple[Tuple[int, ...], Tuple[int, ...]]]] = None, enable_dynamic_routing: bool = True, pid_kp: float = 0.1, pid_ki: float = 0.01, pid_kd: float = 0.05, hysteresis_threshold: float = 0.15, attention_sink_size: int = 4, kv_cache_window_size: int = 1024, quantize_kv_cold: bool = True, kv_quant_bits: int = 4, target_energy_budget: float = 1.0, rope_theta: float = 10000.0, initializer_range: float = 0.02, tie_word_embeddings: bool = False, use_cache: bool = True, bos_token_id: Optional[int] = 1, eos_token_id: Optional[int] = 2, pad_token_id: Optional[int] = None, **kwargs, ): # Resolve legacy parameter aliases if provided in kwargs if "d_model" in kwargs: hidden_size = kwargs.pop("d_model") if "n_layers" in kwargs: num_hidden_layers = kwargs.pop("n_layers") if "n_heads" in kwargs: num_attention_heads = kwargs.pop("n_heads") if "n_kv_heads" in kwargs: num_key_value_heads = kwargs.pop("n_kv_heads") if "max_seq_len" in kwargs: max_position_embeddings = kwargs.pop("max_seq_len") if "ff_multiplier" in kwargs and intermediate_size == 5632: ff_mult = kwargs.pop("ff_multiplier") intermediate_size = hidden_size * ff_mult if "tt_rank" in kwargs and tt_ranks is None: r = kwargs.pop("tt_rank") tt_ranks = [1, r, r, 1] self.vocab_size = vocab_size self.hidden_size = hidden_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads # Key-Value heads: default to GQA 4 if heads >= 4, else match num_attention_heads if num_key_value_heads is None: num_key_value_heads = min(4, num_attention_heads) elif num_key_value_heads > num_attention_heads: num_key_value_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.intermediate_size = intermediate_size self.max_position_embeddings = max_position_embeddings self.rms_norm_eps = rms_norm_eps # Tensor-Train parameters self.tt_ranks = tt_ranks if tt_ranks is not None else [1, 16, 16, 1] self.tt_shapes = tt_shapes # Closed-loop PID & Dynamic Routing parameters self.enable_dynamic_routing = enable_dynamic_routing self.pid_kp = pid_kp self.pid_ki = pid_ki self.pid_kd = pid_kd self.hysteresis_threshold = hysteresis_threshold self.target_energy_budget = target_energy_budget # Adaptive KV Cache parameters self.attention_sink_size = attention_sink_size self.kv_cache_window_size = kv_cache_window_size self.quantize_kv_cold = quantize_kv_cold self.kv_quant_bits = kv_quant_bits # Attention & Position parameters self.rope_theta = rope_theta self.initializer_range = initializer_range self.use_cache = use_cache self.bos_token_id = bos_token_id self.eos_token_id = eos_token_id self.pad_token_id = pad_token_id 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, ) self._validate_config() def _validate_config(self): if self.hidden_size % self.num_attention_heads != 0: raise ValueError( f"hidden_size ({self.hidden_size}) must be divisible by num_attention_heads ({self.num_attention_heads})" ) if self.num_attention_heads % self.num_key_value_heads != 0: raise ValueError( f"num_attention_heads ({self.num_attention_heads}) must be divisible by num_key_value_heads ({self.num_key_value_heads})" ) if len(self.tt_ranks) < 3: raise ValueError(f"tt_ranks must contain at least 3 elements [1, ..., 1], got {self.tt_ranks}") if self.tt_ranks[0] != 1 or self.tt_ranks[-1] != 1: raise ValueError(f"Boundary tt_ranks must be 1, got {self.tt_ranks}") if self.kv_quant_bits not in (4, 8): raise ValueError(f"kv_quant_bits must be either 4 or 8, got {self.kv_quant_bits}") @property def head_dim(self) -> int: return self.hidden_size // self.num_attention_heads @property def num_key_value_groups(self) -> int: return self.num_attention_heads // self.num_key_value_heads @property def d_model(self) -> int: return self.hidden_size @property def n_layers(self) -> int: return self.num_hidden_layers @property def n_heads(self) -> int: return self.num_attention_heads @property def n_kv_heads(self) -> int: return self.num_key_value_heads @property def max_seq_len(self) -> int: return self.max_position_embeddings def get_tt_shapes(self, in_features: int, out_features: int) -> Tuple[Tuple[int, ...], Tuple[int, ...]]: """Compute or retrieve input and output tensor modes for TTLinear.""" num_cores = len(self.tt_ranks) - 1 in_shape = factorize_dim(in_features, num_cores=num_cores) out_shape = factorize_dim(out_features, num_cores=num_cores) return in_shape, out_shape # Register QTensorFormerConfig with Hugging Face AutoConfig AutoConfig.register("qtensorformer", QTensorFormerConfig)