Text Generation
Transformers
Safetensors
PyTorch
English
qtensorformer
q-tensorformer
tensor-train
tensor-networks
adaptive-computation
kv-cache
kv-cache-quantization
attention-sinks
pid-control
edge-ai
edge-inference
on-device-ai
model-compression
grouped-query-attention
rotary-position-embeddings
custom_code
| """ | |
| 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}") | |
| def head_dim(self) -> int: | |
| return self.hidden_size // self.num_attention_heads | |
| def num_key_value_groups(self) -> int: | |
| return self.num_attention_heads // self.num_key_value_heads | |
| def d_model(self) -> int: | |
| return self.hidden_size | |
| def n_layers(self) -> int: | |
| return self.num_hidden_layers | |
| def n_heads(self) -> int: | |
| return self.num_attention_heads | |
| def n_kv_heads(self) -> int: | |
| return self.num_key_value_heads | |
| 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) | |