from __future__ import annotations import math import torch import torch.nn as nn import torch.nn.functional as F from transformers.cache_utils import Cache, DynamicCache from transformers import GenerationMixin from transformers import PreTrainedModel from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast try: from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS except ImportError: # pragma: no cover - compatibility with older Transformers. ALL_ATTENTION_FUNCTIONS = None from .configuration_talkie import TalkieConfig def eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor | None, dropout: float = 0.0, scaling: float | None = None, is_causal: bool | None = None, **kwargs, ) -> tuple[torch.Tensor, None]: del kwargs is_causal = is_causal if is_causal is not None else getattr(module, "is_causal", True) output = F.scaled_dot_product_attention( query, key, value, attn_mask=attention_mask, dropout_p=dropout, scale=scaling, is_causal=is_causal and attention_mask is None, ) return output.transpose(1, 2).contiguous(), None def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: d = x.shape[3] // 2 x1 = x[..., :d] x2 = x[..., d:] y1 = x1 * cos + x2 * sin y2 = x1 * (-sin) + x2 * cos return torch.cat([y1, y2], 3).type_as(x) class HeadGain(nn.Module): def __init__(self, n_head: int): super().__init__() self.head_g = nn.Parameter(torch.ones([n_head])) def forward(self, x: torch.Tensor) -> torch.Tensor: return x * self.head_g.type_as(x).view(1, 1, -1, 1) class WeightGain(nn.Module): def __init__(self): super().__init__() self.w_g = nn.Parameter(torch.ones(1)) def forward(self, w: torch.Tensor) -> torch.Tensor: return w * self.w_g.type_as(w) class ActGain(nn.Module): def __init__(self, init_value: float): super().__init__() self.a_g = nn.Parameter(torch.ones(1) * init_value) def forward(self, x: torch.Tensor) -> torch.Tensor: return x * self.a_g.type_as(x) class CausalSelfAttention(nn.Module): is_causal = True def __init__(self, config: TalkieConfig, layer_idx: int): super().__init__() self.config = config self.layer_idx = layer_idx self.n_head = config.n_head self.head_dim = config.head_dim n_state = config.n_embd self.attn_query = nn.Linear(n_state, n_state, bias=False) self.attn_key = nn.Linear(n_state, n_state, bias=False) self.attn_value = nn.Linear(n_state, n_state, bias=False) self.attn_resid = nn.Linear(n_state, n_state, bias=False) self.head_gain = HeadGain(config.n_head) def forward( self, x: torch.Tensor, cos_sin: tuple[torch.Tensor, torch.Tensor], attention_mask: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: bsz, seq_len, _ = x.size() q = self.attn_query(x).view(bsz, seq_len, self.n_head, self.head_dim) k = self.attn_key(x).view(bsz, seq_len, self.n_head, self.head_dim) v = self.attn_value(x).view(bsz, seq_len, self.n_head, self.head_dim) cos, sin = cos_sin q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin) q, k = F.rms_norm(q, (q.size(-1),)), F.rms_norm(k, (k.size(-1),)) q = self.head_gain(q) key_states = k.transpose(1, 2) value_states = v.transpose(1, 2) if kwargs.get("past_key_values") is not None: key_states, value_states = kwargs["past_key_values"].update( key_states, value_states, self.layer_idx ) if ALL_ATTENTION_FUNCTIONS is None: attention_interface = eager_attention_forward elif hasattr(ALL_ATTENTION_FUNCTIONS, "get_interface"): attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) else: # pragma: no cover - compatibility with older Transformers. attention_interface = ALL_ATTENTION_FUNCTIONS.get( self.config._attn_implementation, eager_attention_forward ) is_causal = attention_mask is None and key_states.shape[-2] == q.shape[1] y, _ = attention_interface( self, q.transpose(1, 2), key_states, value_states, attention_mask, is_causal=is_causal, **kwargs, ) y = y.contiguous().view_as(x) return self.attn_resid(y) class MLP(nn.Module): def __init__(self, config: TalkieConfig): super().__init__() n_state = config.n_embd n_mlp = int(round(((8 / 3) * n_state) / 128) * 128) self.mlp_gate = nn.Linear(n_state, n_mlp, bias=False) self.mlp_linear = nn.Linear(n_state, n_mlp, bias=False) self.mlp_resid = nn.Linear(n_mlp, n_state, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: x = F.silu(self.mlp_gate(x)) * self.mlp_linear(x) return self.mlp_resid(x) class Block(nn.Module): def __init__(self, config: TalkieConfig, layer_idx: int): super().__init__() self.attn = CausalSelfAttention(config, layer_idx) self.attn_gain = ActGain((2 * config.n_layer) ** -0.5) self.mlp = MLP(config) self.mlp_gain = ActGain((2 * config.n_layer) ** -0.5) self.embed_skip = ActGain(0.0) def forward( self, e_x: torch.Tensor, x: torch.Tensor, cos_sin: tuple[torch.Tensor, torch.Tensor], attention_mask: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: x = x + self.attn_gain( self.attn(F.rms_norm(x, (x.shape[-1],)), cos_sin, attention_mask, **kwargs) ) x = x + self.mlp_gain(self.mlp(F.rms_norm(x, (x.shape[-1],)))) x = x + self.embed_skip(e_x) return x class TalkiePreTrainedModel(PreTrainedModel): config_class = TalkieConfig base_model_prefix = "" supports_gradient_checkpointing = True _supports_sdpa = True _supports_attention_backend = True _no_split_modules = ["Block"] _tied_weights_keys = None def _init_weights(self, module: nn.Module) -> None: return class TalkieModel(TalkiePreTrainedModel, GenerationMixin): def __init__(self, config: TalkieConfig): super().__init__(config) self.embed = nn.Embedding(config.vocab_size, config.n_embd) self.blocks = nn.ModuleList([Block(config, i) for i in range(config.n_layer)]) self.gradient_checkpointing = False cos, sin = self._precompute_rotary_embeddings(config.max_position_embeddings) self.register_buffer("cos", cos, persistent=False) self.register_buffer("sin", sin, persistent=False) self._rotary_initialized = cos.device.type != "meta" self.post_init() def _precompute_rotary_embeddings( self, seq_len: int, head_dim: int | None = None, base: int | float | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: device = self.embed.weight.device if hasattr(self, "embed") else "cpu" head_dim = head_dim if head_dim is not None else self.config.head_dim base = base if base is not None else self.config.rope_base inv_freq, attention_factor = self._rotary_inv_freq(seq_len, head_dim, float(base), device) t = torch.arange(seq_len, dtype=torch.float32, device=device) freqs = torch.outer(t, inv_freq) cos, sin = freqs.cos(), freqs.sin() if attention_factor != 1.0: cos = cos * attention_factor sin = sin * attention_factor cos, sin = cos.bfloat16(), sin.bfloat16() cos, sin = cos[None, :, None, :], sin[None, :, None, :] return cos, sin def _rotary_inv_freq( self, seq_len: int, head_dim: int, base: float, device: torch.device | str, ) -> tuple[torch.Tensor, float]: scaling = self.config.rope_scaling rope_type = scaling.get("rope_type") if scaling else None if rope_type in (None, "default"): return self._default_rotary_inv_freq(head_dim, base, device), 1.0 if rope_type == "linear": inv_freq = self._default_rotary_inv_freq(head_dim, base, device) return inv_freq / float(scaling["factor"]), 1.0 if rope_type == "dynamic": return self._dynamic_rotary_inv_freq(seq_len, head_dim, base, device, scaling), 1.0 if rope_type == "yarn": return self._yarn_rotary_inv_freq(head_dim, base, device, scaling) raise ValueError(f"unsupported rope_scaling type {rope_type!r}") @staticmethod def _default_rotary_inv_freq( head_dim: int, base: float, device: torch.device | str ) -> torch.Tensor: channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) return 1.0 / (base ** (channel_range / head_dim)) def _original_max_position_embeddings(self, scaling: dict | None) -> int: if scaling and "original_max_position_embeddings" in scaling: return int(scaling["original_max_position_embeddings"]) return int(self.config.max_position_embeddings) def _dynamic_rotary_inv_freq( self, seq_len: int, head_dim: int, base: float, device: torch.device | str, scaling: dict, ) -> torch.Tensor: original_max_position_embeddings = self._original_max_position_embeddings(scaling) scaled_seq_len = max(seq_len, original_max_position_embeddings) factor = float(scaling["factor"]) base = base * ( (factor * scaled_seq_len / original_max_position_embeddings) - (factor - 1.0) ) ** (head_dim / (head_dim - 2.0)) return self._default_rotary_inv_freq(head_dim, base, device) def _yarn_rotary_inv_freq( self, head_dim: int, base: float, device: torch.device | str, scaling: dict, ) -> tuple[torch.Tensor, float]: factor = float(scaling["factor"]) original_max_position_embeddings = self._original_max_position_embeddings(scaling) beta_fast = float(scaling.get("beta_fast", 32.0)) beta_slow = float(scaling.get("beta_slow", 1.0)) attention_factor = scaling.get("attention_factor") if attention_factor is None: attention_factor = 1.0 if factor <= 1.0 else 0.1 * math.log(factor) + 1.0 channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) pos_freqs = base ** (channel_range / head_dim) inv_freq_extrapolation = 1.0 / pos_freqs inv_freq_interpolation = 1.0 / (factor * pos_freqs) low, high = self._yarn_correction_range( beta_fast, beta_slow, head_dim, base, original_max_position_embeddings, truncate=bool(scaling.get("truncate", True)), ) ramp = self._yarn_linear_ramp(low, high, head_dim // 2, device) extrapolation_factor = 1.0 - ramp inv_freq = ( inv_freq_interpolation * (1.0 - extrapolation_factor) + inv_freq_extrapolation * extrapolation_factor ) return inv_freq, float(attention_factor) @staticmethod def _yarn_correction_range( low_rot: float, high_rot: float, head_dim: int, base: float, original_max_position_embeddings: int, truncate: bool, ) -> tuple[float, float]: def correction_dim(num_rotations: float) -> float: return ( head_dim * math.log(original_max_position_embeddings / (num_rotations * 2.0 * math.pi)) / (2.0 * math.log(base)) ) low = correction_dim(low_rot) high = correction_dim(high_rot) if truncate: low = math.floor(low) high = math.ceil(high) return max(low, 0.0), min(high, float(head_dim - 1)) @staticmethod def _yarn_linear_ramp( low: float, high: float, dim: int, device: torch.device | str, ) -> torch.Tensor: if low == high: high += 0.001 ramp = (torch.arange(dim, dtype=torch.float32, device=device) - low) / (high - low) return torch.clamp(ramp, 0.0, 1.0) def _ensure_rotary_embeddings(self, seq_len: int) -> None: device = self.embed.weight.device needs_init = ( not self._rotary_initialized or self.cos.device != device or self.sin.device != device or self.cos.shape[1] < seq_len ) if needs_init: max_seq_len = max(seq_len, self.config.max_position_embeddings) cos, sin = self._precompute_rotary_embeddings(max_seq_len) self.cos = cos.to(device=device) self.sin = sin.to(device=device) self._rotary_initialized = True def reset_rotary_embeddings(self) -> None: self._rotary_initialized = False def get_input_embeddings(self) -> nn.Embedding: return self.embed def set_input_embeddings(self, value: nn.Embedding) -> None: self.embed = value def _position_ids( self, input_ids: torch.LongTensor, position_ids: torch.LongTensor | None = None, cache_position: torch.LongTensor | None = None, past_key_values: Cache | None = None, ) -> torch.LongTensor: batch_size, seq_len = input_ids.shape if position_ids is not None: if position_ids.dim() == 1: position_ids = position_ids.unsqueeze(0) return position_ids.to(device=input_ids.device, dtype=torch.long) if cache_position is not None: if cache_position.dim() == 1: cache_position = cache_position.unsqueeze(0) if cache_position.shape[0] == 1 and batch_size != 1: cache_position = cache_position.expand(batch_size, -1) return cache_position.to(device=input_ids.device, dtype=torch.long) past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0 position_ids = torch.arange(seq_len, device=input_ids.device, dtype=torch.long) + past_seen return position_ids.unsqueeze(0).expand(batch_size, -1) def _attention_mask( self, attention_mask: torch.Tensor | None, input_ids: torch.Tensor, position_ids: torch.Tensor, past_key_values: Cache | None, dtype: torch.dtype, ) -> torch.Tensor | None: if attention_mask is not None and attention_mask.dim() >= 4: return attention_mask batch_size, query_length = input_ids.shape past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0 if attention_mask is not None and attention_mask.dim() != 2: return attention_mask if attention_mask is None and past_seen == 0: return None key_length = past_seen + query_length if attention_mask is not None: if attention_mask.shape[-1] == query_length and past_seen: prefix = torch.ones( attention_mask.shape[0], past_seen, dtype=attention_mask.dtype, device=attention_mask.device, ) attention_mask = torch.cat([prefix, attention_mask], dim=-1) key_length = attention_mask.shape[-1] key_positions = torch.arange(key_length, device=input_ids.device, dtype=torch.long) future_mask = key_positions.view(1, 1, 1, key_length) > position_ids.view( batch_size, 1, query_length, 1 ) if attention_mask is not None: padding_mask = attention_mask[:, None, None, :].to(device=input_ids.device) == 0 mask = future_mask | padding_mask else: mask = future_mask min_value = torch.finfo(dtype).min causal_mask = torch.zeros( batch_size, 1, query_length, key_length, dtype=dtype, device=input_ids.device ) return causal_mask.masked_fill(mask, min_value) def forward( self, input_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, use_cache: bool | None = None, return_dict: bool | None = None, **kwargs, ) -> BaseModelOutputWithPast | tuple[torch.Tensor, ...]: cache_position = kwargs.pop("cache_position", None) if input_ids is None and inputs_embeds is None: raise ValueError("input_ids or inputs_embeds is required") if input_ids is not None and inputs_embeds is not None: raise ValueError("provide only one of input_ids or inputs_embeds") if input_ids is None: input_ids = torch.empty( inputs_embeds.shape[:2], dtype=torch.long, device=inputs_embeds.device, ) use_cache = use_cache if use_cache is not None else self.config.use_cache if self.gradient_checkpointing and self.training: use_cache = False if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) position_ids = self._position_ids(input_ids, position_ids, cache_position, past_key_values) # Keep graph capture free of CUDA tensor -> Python scalar syncs. The # configured context length is the static serving/training contract. self._ensure_rotary_embeddings(int(self.config.max_position_embeddings)) cos = self.cos[0, position_ids, :, :] sin = self.sin[0, position_ids, :, :] cos_sin = cos, sin x = inputs_embeds if inputs_embeds is not None else self.embed(input_ids) x = F.rms_norm(x, (x.shape[-1],)) attention_mask = self._attention_mask(attention_mask, input_ids, position_ids, past_key_values, x.dtype) e_x = x for block in self.blocks: if self.gradient_checkpointing and self.training: def custom_forward( e_x: torch.Tensor, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, attention_mask: torch.Tensor | None, block: Block = block, ) -> torch.Tensor: return block(e_x, x, (cos, sin), attention_mask=attention_mask) x = self._gradient_checkpointing_func( custom_forward, e_x, x, cos, sin, attention_mask, ) else: x = block( e_x, x, cos_sin, attention_mask=attention_mask, past_key_values=past_key_values if use_cache else None, **kwargs, ) x = F.rms_norm(x, (x.shape[-1],)) past_key_values = past_key_values if use_cache else None use_return_dict = return_dict if return_dict is not None else self.config.use_return_dict if use_return_dict: return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=past_key_values) output = (x,) return output + ((past_key_values,) if past_key_values is not None else ()) class TalkieForCausalLM(TalkieModel): _tied_weights_keys = None def __init__(self, config: TalkieConfig): super().__init__(config) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) self.post_init() def get_output_embeddings(self) -> nn.Linear: return self.lm_head def set_output_embeddings(self, value: nn.Linear) -> None: self.lm_head = value def _chunked_lm_loss( self, hidden_states: torch.Tensor, labels: torch.Tensor, chunk_size: int, ) -> torch.Tensor: if chunk_size <= 0: raise ValueError("chunk_size must be positive") total_loss = hidden_states.new_zeros((), dtype=torch.float32) total_tokens = hidden_states.new_zeros((), dtype=torch.float32) for start in range(0, hidden_states.shape[1], chunk_size): end = min(start + chunk_size, hidden_states.shape[1]) logits = self.lm_head(hidden_states[:, start:end, :]).float() if self.config.logit_scale != 1.0: logits = logits * self.config.logit_scale chunk_labels = labels[:, start:end].contiguous() total_loss = total_loss + F.cross_entropy( logits.reshape(-1, logits.size(-1)), chunk_labels.reshape(-1), ignore_index=-100, reduction="sum", ) total_tokens = total_tokens + (chunk_labels != -100).sum(dtype=torch.float32) return total_loss / total_tokens.clamp_min(1.0) def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, return_dict: bool | None = None, past_key_values: Cache | None = None, use_cache: bool | None = None, position_ids: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, loss_chunk_size: int = 0, return_logits: bool = True, **kwargs, ) -> CausalLMOutputWithPast | tuple[torch.Tensor, ...]: if input_ids is None and inputs_embeds is None: raise ValueError("input_ids or inputs_embeds is required") cache_position = kwargs.pop("cache_position", None) outputs = super().forward( input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, cache_position=cache_position, return_dict=True, **kwargs, ) hidden_states = outputs.last_hidden_state loss = None logits = None if labels is not None and loss_chunk_size > 0: loss = self._chunked_lm_loss( hidden_states[:, :-1, :], labels[:, 1:], loss_chunk_size, ) if return_logits: slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]).float() if self.config.logit_scale != 1.0: logits = logits * self.config.logit_scale if labels is not None and loss is None: if logits is None: raise ValueError("return_logits must be true when loss_chunk_size is not used") shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss = F.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100, ) use_return_dict = return_dict if return_dict is not None else self.config.use_return_dict if use_return_dict: return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, ) output = (logits,) if outputs.past_key_values is not None: output += (outputs.past_key_values,) return ((loss,) + output) if loss is not None else output __all__ = ["TalkieConfig", "TalkieForCausalLM", "TalkieModel"]