"""HF-compatible inference wrapper for Slayer's trained BalancedLM weights.""" import math import torch from torch import nn from torch.nn import functional as F from transformers import PretrainedConfig, PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithPast class SlayerBalancedConfig(PretrainedConfig): model_type = "slayer_balanced" def __init__(self, vocab_size=24576, hidden_size=576, intermediate_size=2496, num_hidden_layers=26, num_attention_heads=9, num_key_value_heads=3, head_dim=64, max_position_embeddings=2048, rms_norm_eps=1e-6, rope_parameters=None, tie_word_embeddings=True, **kwargs): super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) self.vocab_size = vocab_size self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.head_dim = head_dim self.max_position_embeddings = max_position_embeddings self.rms_norm_eps = rms_norm_eps self.rope_parameters = rope_parameters or {"rope_theta": 10000.0, "rope_type": "default"} self.use_cache = False class RMSNorm(nn.Module): def __init__(self, width, eps): super().__init__() self.weight = nn.Parameter(torch.ones(width)) self.eps = eps def forward(self, x): dtype = x.dtype x = x.float() return self.weight * (x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps)).to(dtype) class ValueProjection(nn.Linear): def __init__(self, width, out, first): super().__init__(width, out, bias=False) if not first: self.vr_lambda = nn.Parameter(torch.zeros(1)) class SlayerAttention(nn.Module): def __init__(self, c, index): super().__init__() h = c.hidden_size self.heads, self.kv, self.d = c.num_attention_heads, c.num_key_value_heads, c.head_dim self.first = index == 0 self.q_proj = nn.Linear(h, self.heads * self.d, bias=False) self.k_proj = nn.Linear(h, self.kv * self.d, bias=False) self.v_proj = ValueProjection(h, self.kv * self.d, self.first) self.o_proj = nn.Linear(self.heads * self.d, h, bias=False) self.q_norm, self.k_norm = RMSNorm(self.d, c.rms_norm_eps), RMSNorm(self.d, c.rms_norm_eps) def forward(self, x, cos, sin, v0, attention_mask=None, output_attentions=False): b, t, _ = x.shape q = self.q_norm(self.q_proj(x).view(b, t, self.heads, self.d)).transpose(1, 2) k = self.k_norm(self.k_proj(x).view(b, t, self.kv, self.d)).transpose(1, 2) v = self.v_proj(x) if self.first: v0 = v else: v = v + self.v_proj.vr_lambda * v0 v = v.view(b, t, self.kv, self.d).transpose(1, 2) if self.heads != self.kv: rep = self.heads // self.kv k = k.repeat_interleave(rep, dim=1) v = v.repeat_interleave(rep, dim=1) def rotate(z): return torch.cat((-z[..., self.d // 2:], z[..., :self.d // 2]), dim=-1) q, k = q * cos + rotate(q) * sin, k * cos + rotate(k) * sin probs = None if output_attentions: scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d) causal = torch.ones((t, t), dtype=torch.bool, device=x.device).tril() scores = scores.masked_fill(~causal, torch.finfo(scores.dtype).min) if attention_mask is not None: scores = scores.masked_fill(~attention_mask[:, None, None, :].bool(), torch.finfo(scores.dtype).min) probs = scores.softmax(dim=-1) out = probs @ v else: if attention_mask is None: out = F.scaled_dot_product_attention(q, k, v, is_causal=True) else: causal = torch.ones((t, t), dtype=torch.bool, device=x.device).tril() allowed = causal[None, None] & attention_mask[:, None, None, :].bool() out = F.scaled_dot_product_attention(q, k, v, attn_mask=allowed, is_causal=False) out = self.o_proj(out.transpose(1, 2).contiguous().view(b, t, -1)) return out, v0, probs class SlayerMLP(nn.Module): def __init__(self, c): super().__init__() h, f = c.hidden_size, c.intermediate_size self.gate_proj = nn.Linear(h, f, bias=False) self.up_proj = nn.Linear(h, f, bias=False) self.down_proj = nn.Linear(f, h, bias=False) def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) class SlayerBlock(nn.Module): def __init__(self, c, index): super().__init__() self.self_attn = SlayerAttention(c, index) self.mlp = SlayerMLP(c) self.input_layernorm = RMSNorm(c.hidden_size, c.rms_norm_eps) self.post_attention_layernorm = RMSNorm(c.hidden_size, c.rms_norm_eps) def forward(self, x, cos, sin, v0, mask, output_attentions): attn, v0, probs = self.self_attn(self.input_layernorm(x), cos, sin, v0, mask, output_attentions) x = x + attn return x + self.mlp(self.post_attention_layernorm(x)), v0, probs class SlayerBackbone(nn.Module): def __init__(self, c): super().__init__() self.embed_tokens = nn.Embedding(c.vocab_size, c.hidden_size) self.layers = nn.ModuleList(SlayerBlock(c, i) for i in range(c.num_hidden_layers)) self.norm = RMSNorm(c.hidden_size, c.rms_norm_eps) class SlayerForCausalLM(PreTrainedModel): config_class = SlayerBalancedConfig base_model_prefix = "model" _no_split_modules = ["SlayerBlock"] _supports_sdpa = True def __init__(self, config): super().__init__(config) self.model = SlayerBackbone(config) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.post_init() def get_input_embeddings(self): return self.model.embed_tokens def set_input_embeddings(self, value): self.model.embed_tokens = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, value): self.lm_head = value def forward(self, input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None, labels=None, output_attentions=None, use_cache=None, return_dict=None, **kwargs): if inputs_embeds is None: inputs_embeds = self.model.embed_tokens(input_ids) _, t, _ = inputs_embeds.shape d = self.config.head_dim theta = self.config.rope_parameters.get("rope_theta", 10000.0) inv = 1.0 / theta ** (torch.arange(0, d, 2, device=inputs_embeds.device, dtype=torch.float32) / d) if position_ids is None: if attention_mask is None: position_ids = torch.arange(t, device=inputs_embeds.device)[None, :].expand(inputs_embeds.shape[0], -1) else: position_ids = (attention_mask.long().cumsum(-1) - 1).clamp_min(0) phase = position_ids.to(torch.float32).unsqueeze(-1) * inv phase = torch.cat((phase, phase), dim=-1) cos = phase.cos()[:, None].to(inputs_embeds.dtype) sin = phase.sin()[:, None].to(inputs_embeds.dtype) x, v0 = inputs_embeds, None attentions = [] if output_attentions else None for layer in self.model.layers: x, v0, attn = layer(x, cos, sin, v0, attention_mask, bool(output_attentions)) if attentions is not None: attentions.append(attn) logits = self.lm_head(self.model.norm(x)) loss = None if labels is not None: loss = F.cross_entropy(logits[:, :-1].float().reshape(-1, logits.size(-1)), labels[:, 1:].reshape(-1)) return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=None, attentions=tuple(attentions) if attentions is not None else None, )