| """Portable ARLoop inference: PyTorch SDPA and FP32 attention-residual routing. |
| |
| Adapted from ARLoop's model and Parcae building blocks; see LICENSE-MIT. |
| 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_arloop import ARLoopConfig |
|
|
|
|
| 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): |
| |
| 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).to(v.dtype) |
| |
| 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: |
| |
| 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 ARLoopBlock(nn.Module): |
| def __init__(self, config, layer): |
| super().__init__() |
| self.attn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) |
| self.attn = Attention(config, layer) |
| self.mlp_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) |
| self.mlp = MLP(config) |
|
|
|
|
| class ARLoopForCausalLM(PreTrainedModel, GenerationMixin): |
| config_class = ARLoopConfig |
| base_model_prefix = "" |
| main_input_name = "input_ids" |
| _no_split_modules = ["ARLoopBlock"] |
| _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(ARLoopBlock(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.attnres_queries = nn.Parameter(torch.zeros(2 * config.num_hidden_layers, config.hidden_size)) |
| self.attnres_norms = nn.ModuleList(RMSNorm(config.hidden_size, config.rms_norm_eps) |
| for _ in range(2 * config.num_hidden_layers)) |
| self.final_attnres_query = nn.Parameter(torch.zeros(config.hidden_size)) |
| self.final_attnres_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) |
| 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 _route(self, query, norm, blocks, partial, output_norm=None): |
| values = blocks + [partial] |
| if len(values) > self.config.sw_blocks + 1: |
| values = [values[0], *values[-self.config.sw_blocks:]] |
| |
| |
| with torch.autocast(device_type=partial.device.type, enabled=False): |
| stacked = torch.stack(values, dim=-2).float() |
| scores = (stacked * (query.float() * norm.weight.float())).sum(-1) |
| scores = scores * torch.rsqrt(stacked.square().mean(-1) + norm.eps) |
| weights = scores.softmax(dim=-1) |
| y = torch.zeros_like(partial, dtype=torch.float32) |
| for i, value in enumerate(values): |
| y = y + weights[..., i, None] * value.float() |
| if output_norm is not None: |
| y = y * torch.rsqrt(y.square().mean(-1, keepdim=True) + norm.eps) |
| y = y * output_norm.weight.float() |
| return y.to(partial.dtype) |
|
|
| 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() |
| |
| with torch.autocast(device_type=input_ids.device.type, dtype=torch.bfloat16): |
| embedded = (self.transformer.wte(input_ids) * self.config.embedding_scale).to(torch.bfloat16) |
| ve = self.shared_value_embedding(input_ids) |
| blocks, partial = [embedded], embedded |
| states = [embedded] 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 = self._route(self.attnres_queries[2 * effective], |
| self.attnres_norms[2 * effective], blocks, partial, block.attn_norm) |
| if effective > 0 and effective % (self.config.attnres_block_size // 2) == 0: |
| blocks.append(partial) |
| partial = None |
| a = block.attn(h, position_ids, mask, ve, cache, effective) |
| partial = a if partial is None else partial + a |
| h = self._route(self.attnres_queries[2 * effective + 1], |
| self.attnres_norms[2 * effective + 1], blocks, partial, block.mlp_norm) |
| partial = partial + block.mlp(h) |
| if states is not None: |
| states.append(partial) |
| h = self._route(self.final_attnres_query, self.final_attnres_norm, blocks, partial) |
| |
| 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) |
|
|
