"""Portable fixed-loop PreNorm baseline inference with PyTorch SDPA. Adapted from ARLoop's model and Parcae building blocks; see NOTICE. No training repository, FLA, Triton, or flash-attn installation is needed. """ import torch from torch import nn from torch.nn import functional as F from transformers import PreTrainedModel, GenerationMixin from transformers.cache_utils import DynamicCache from transformers.modeling_outputs import CausalLMOutputWithPast from .configuration_baseline import ARLoopBaselineConfig class RMSNorm(nn.Module): def __init__(self, size, eps): super().__init__() self.weight = nn.Parameter(torch.ones(size)) self.eps = eps def forward(self, x): # Preserve the native final norm's rounding and weight promotion. y = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + self.eps) return y.to(x.dtype) * self.weight class Attention(nn.Module): def __init__(self, config, layer): super().__init__() self.config = config dim = config.hidden_size kv_dim = config.num_key_value_heads * config.head_dim self.c_q = nn.Linear(dim, dim, bias=False) self.c_k = nn.Linear(dim, kv_dim, bias=False) self.c_v = nn.Linear(dim, kv_dim, bias=False) self.c_proj = nn.Linear(dim, dim, bias=False) self.ve_gate = (nn.Linear(32, config.num_key_value_heads, bias=False) if layer % 2 == (config.physical_layers - 1) % 2 else None) def forward(self, x, positions, mask, value_embedding, cache, effective_layer): c = self.config batch, length, _ = x.shape q = self.c_q(x).view(batch, length, c.num_attention_heads, c.head_dim) k = self.c_k(x).view(batch, length, c.num_key_value_heads, c.head_dim) v = self.c_v(x).view_as(k) if self.ve_gate is not None: gate = 2 * torch.sigmoid(self.ve_gate(x[..., :32])) v = v + gate.unsqueeze(-1) * value_embedding.view_as(v) # Adjacent-pair rotary embedding, unlike Llama's split-half convention. freq = 1.0 / (c.rope_theta ** (torch.arange(0, c.head_dim, 2, device=x.device, dtype=torch.float32) / c.head_dim)) angles = positions.float().unsqueeze(-1) / c.rope_condense_ratio * freq cos, sin = angles.cos().unsqueeze(2), angles.sin().unsqueeze(2) def rotate(tensor): pairs = tensor.float().unflatten(-1, (-1, 2)) return torch.stack((pairs[..., 0] * cos - pairs[..., 1] * sin, pairs[..., 1] * cos + pairs[..., 0] * sin), -1).flatten(-2).to(tensor.dtype) q, k = rotate(q), rotate(k) if c.qk_norm: # Native Q/K normalization uses dtype-dependent default epsilon. q = F.rms_norm(q, (c.head_dim,)) k = F.rms_norm(k, (c.head_dim,)) q, k, v = (t.transpose(1, 2) for t in (q, k, v)) if cache is not None: k, v = cache.update(k, v, effective_layer) causal = mask is None and q.shape[-2] == k.shape[-2] if mask is None and q.shape[-2] != k.shape[-2] and q.shape[-2] != 1: rows = torch.arange(k.shape[-2] - q.shape[-2], k.shape[-2], device=x.device) mask = torch.arange(k.shape[-2], device=x.device)[None, :] <= rows[:, None] y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=causal, enable_gqa=c.num_attention_heads != c.num_key_value_heads, dropout_p=0.0) return self.c_proj(y.transpose(1, 2).contiguous().view(batch, length, c.hidden_size)) class MLP(nn.Module): def __init__(self, config): super().__init__() self.fc = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False) self.proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) def forward(self, x): gate, value = self.fc(x).chunk(2, dim=-1) return self.proj(F.silu(gate) * value) class PreNormBlock(nn.Module): def __init__(self, config, layer): super().__init__() self.norm_1 = RMSNorm(config.hidden_size, config.rms_norm_eps) self.attn = Attention(config, layer) self.norm_2 = RMSNorm(config.hidden_size, config.rms_norm_eps) self.mlp = MLP(config) class ARLoopBaselineForCausalLM(PreTrainedModel, GenerationMixin): config_class = ARLoopBaselineConfig base_model_prefix = "" main_input_name = "input_ids" _no_split_modules = ["PreNormBlock"] _supports_cache_class = True def __init__(self, config): super().__init__(config) self.transformer = nn.ModuleDict({ "wte": nn.Embedding(config.vocab_size, config.hidden_size), "h": nn.ModuleList(PreNormBlock(config, i) for i in range(config.physical_layers)), "ln_f": RMSNorm(config.hidden_size, config.rms_norm_eps), }) self.shared_value_embedding = nn.Embedding(config.vocab_size, config.num_key_value_heads * config.head_dim) self.post_init() def _init_weights(self, module): if isinstance(module, (nn.Linear, nn.Embedding)): nn.init.normal_(module.weight, std=0.02) elif isinstance(module, RMSNorm): nn.init.ones_(module.weight) def get_input_embeddings(self): return self.transformer.wte def set_input_embeddings(self, value): self.transformer.wte = value def forward(self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, labels=None, use_cache=None, return_dict=None, output_hidden_states=False, output_attentions=False, cache_position=None, logits_to_keep=0, **kwargs): if input_ids is None or inputs_embeds is not None: raise ValueError("ARLoop requires input_ids for token-dependent value embeddings") if output_attentions: raise ValueError("output_attentions is not implemented") use_cache = self.config.use_cache if use_cache is None else use_cache if labels is not None: use_cache = False if past_key_values is not None and not use_cache: raise ValueError("past_key_values requires use_cache=True") cache = past_key_values if use_cache else None if use_cache and cache is None: cache = DynamicCache() past_length = cache.get_seq_length() if cache is not None else 0 batch, length = input_ids.shape total = past_length + length if total > self.config.max_position_embeddings: raise ValueError("Input and cache exceed max_position_embeddings") if position_ids is None: if attention_mask is None: position_ids = torch.arange(past_length, total, device=input_ids.device)[None, :] else: position_ids = (attention_mask.long().cumsum(-1) - 1).clamp_min(0)[:, -length:] key_indices = torch.arange(total, device=input_ids.device) query_indices = torch.arange(past_length, total, device=input_ids.device) mask = None if attention_mask is not None: if attention_mask.shape != (batch, total): raise ValueError("attention_mask must cover both cached and new tokens") if not bool(attention_mask.all()): mask = (key_indices[None, :] <= query_indices[:, None])[None, None, :, :] mask = mask & attention_mask[:, None, None, :].bool() # Residual/value embeddings follow the loading dtype; Linear ops use BF16 autocast. with torch.autocast(device_type=input_ids.device.type, dtype=torch.bfloat16): h = self.transformer.wte(input_ids) * self.config.embedding_scale ve = self.shared_value_embedding(input_ids) states = [h] if output_hidden_states else None for loop in range(self.config.loop_count): for i, block in enumerate(self.transformer.h): effective = loop * self.config.physical_layers + i h = h + block.attn(block.norm_1(h), position_ids, mask, ve, cache, effective) h = h + block.mlp(block.norm_2(h)) if states is not None: states.append(h) # Native RMSNorm disables autocast around the final normalization. with torch.autocast(device_type=input_ids.device.type, enabled=False): h = self.transformer.ln_f(h) if logits_to_keep and labels is None: h = h[:, -logits_to_keep:, :] logits = F.linear(h, self.transformer.wte.weight).float() if self.config.logit_softcap is not None: cap = self.config.logit_softcap logits = cap * torch.tanh(logits / cap) loss = None if labels is not None: loss = F.cross_entropy(logits[:, :-1].contiguous().view(-1, logits.size(-1)), labels[:, 1:].contiguous().view(-1), ignore_index=-100) result = CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=cache, hidden_states=tuple(states) if states is not None else None) return result if return_dict is not False else result.to_tuple() def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs): if past_key_values is not None: seen = past_key_values.get_seq_length() if seen: input_ids = input_ids[:, seen:] if input_ids.shape[1] > seen else input_ids[:, -1:] return dict(input_ids=input_ids.contiguous(), past_key_values=past_key_values, attention_mask=attention_mask, use_cache=kwargs.get("use_cache", True), logits_to_keep=1)