| |
| |
| |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
| from transformers import PretrainedConfig, PreTrainedModel |
| from transformers.modeling_outputs import CausalLMOutputWithPast |
| try: |
| from transformers.generation import GenerationMixin |
| except ImportError: |
| from transformers.generation.utils import GenerationMixin |
|
|
|
|
| class InfernoConfig(PretrainedConfig): |
| model_type = "inferno" |
|
|
| def __init__(self, vocab_size=32064, hidden_size=640, num_hidden_layers=14, num_attention_heads=8, |
| num_key_value_heads=2, head_dim=80, intermediate_size=1536, rope_theta=10000.0, |
| max_position_embeddings=2048, gate_shift_start=-1, tie_word_embeddings=True, **kwargs): |
| self.vocab_size = vocab_size |
| self.hidden_size = hidden_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.intermediate_size = intermediate_size |
| self.rope_theta = rope_theta |
| self.max_position_embeddings = max_position_embeddings |
| self.gate_shift_start = gate_shift_start |
| super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) |
|
|
|
|
| class RMSNorm(nn.RMSNorm): |
| def forward(self, x): |
| return F.rms_norm(x.float(), self.normalized_shape, self.weight, self.eps).to(x.dtype) |
|
|
|
|
| def rope_tables(length, dim, base): |
| inv = 1. / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) |
| freqs = torch.outer(torch.arange(length, dtype=torch.float32), inv) |
| return freqs.cos(), freqs.sin() |
|
|
|
|
| def apply_rope(x, cos, sin): |
| work = x.float() if x.dtype in (torch.float16, torch.bfloat16) else x |
| a, b = work.chunk(2, -1) |
| return torch.cat((a * cos - b * sin, a * sin + b * cos), -1).to(x.dtype) |
|
|
|
|
| class GQAttention(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.H = c.num_attention_heads; self.K = c.num_key_value_heads; self.hd = c.head_dim |
| self.q_width = self.H * self.hd; self.kv_width = self.K * self.hd |
| self.qkv_proj = nn.Linear(c.hidden_size, self.q_width + 2 * self.kv_width, bias=False) |
| self.o_proj = nn.Linear(self.q_width, c.hidden_size, bias=False) |
|
|
| def forward(self, x, cos, sin, attn_mask): |
| B, T, _ = x.shape |
| q, k, v = self.qkv_proj(x).split((self.q_width, self.kv_width, self.kv_width), -1) |
| q = q.view(B, T, self.H, self.hd).transpose(1, 2) |
| k = k.view(B, T, self.K, self.hd).transpose(1, 2) |
| v = v.view(B, T, self.K, self.hd).transpose(1, 2) |
| q = apply_rope(q, cos, sin); k = apply_rope(k, cos, sin) |
| if attn_mask is None: |
| y = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=self.H != self.K) |
| else: |
| y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, enable_gqa=self.H != self.K) |
| return self.o_proj(y.transpose(1, 2).reshape(B, T, self.q_width)) |
|
|
|
|
| class FeedForward(nn.Module): |
| def __init__(self, c, index): |
| super().__init__() |
| self.shift = c.gate_shift_start >= 0 and index >= c.gate_shift_start |
| self.project = nn.Linear(c.hidden_size, 2 * c.intermediate_size, bias=False) |
| self.down = nn.Linear(c.intermediate_size, c.hidden_size, bias=False) |
|
|
| def forward(self, n, own_gate): |
| g, u = self.project(n).chunk(2, -1) |
| if self.shift: |
| previous = torch.cat((g[:, :1], g[:, :-1]), 1) |
| g = previous if own_gate is None else torch.where(own_gate[..., None], g, previous) |
| return self.down(F.silu(g) * u) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, c, index): |
| super().__init__() |
| self.attn_norm = RMSNorm(c.hidden_size); self.attn = GQAttention(c) |
| self.ffn_norm = RMSNorm(c.hidden_size); self.ffn = FeedForward(c, index) |
|
|
| def forward(self, x, cos, sin, attn_mask, own_gate): |
| x = x + self.attn(self.attn_norm(x), cos, sin, attn_mask) |
| return x + self.ffn(self.ffn_norm(x), own_gate) |
|
|
|
|
| class InfernoForCausalLM(PreTrainedModel, GenerationMixin): |
| config_class = InfernoConfig |
| base_model_prefix = "model" |
| _tied_weights_keys = ["lm_head.weight"] |
| _no_split_modules = ["Block"] |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.embed = nn.Embedding(config.vocab_size, config.hidden_size) |
| self.blocks = nn.ModuleList([Block(config, i) for i in range(config.num_hidden_layers)]) |
| self.ln_f = RMSNorm(config.hidden_size) |
| self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
| self.lm_head.weight = self.embed.weight |
| self._rope = {} |
| self.post_init() |
|
|
| def _init_weights(self, module): |
| if isinstance(module, (nn.Linear, nn.Embedding)): |
| nn.init.normal_(module.weight, std=.02) |
|
|
| def get_input_embeddings(self): return self.embed |
| def set_input_embeddings(self, value): self.embed = value |
| def get_output_embeddings(self): return self.lm_head |
| def set_output_embeddings(self, value): self.lm_head = value |
|
|
| def _tables(self, device): |
| if device not in self._rope: |
| c = self.config |
| cos, sin = rope_tables(c.max_position_embeddings + 1, c.head_dim, c.rope_theta) |
| self._rope[device] = (cos.to(device), sin.to(device)) |
| return self._rope[device] |
|
|
| def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs): |
| |
| limit = self.config.max_position_embeddings |
| inputs = {"input_ids": input_ids[:, -limit:]} |
| if attention_mask is not None: |
| inputs["attention_mask"] = attention_mask[:, -limit:] |
| return inputs |
|
|
| def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None, **kwargs): |
| B, T = input_ids.shape |
| if T > self.config.max_position_embeddings: |
| raise ValueError(f"Sequence length {T} exceeds the trained context {self.config.max_position_embeddings}") |
| cos_t, sin_t = self._tables(input_ids.device) |
| attn_mask = own_gate = None |
| if attention_mask is None or bool(attention_mask.all()): |
| cos, sin = cos_t[:T], sin_t[:T] |
| else: |
| keep = attention_mask.bool() |
| positions = (keep.long().cumsum(-1) - 1).clamp(min=0) |
| cos, sin = cos_t[positions].unsqueeze(1), sin_t[positions].unsqueeze(1) |
| causal = torch.ones(T, T, dtype=torch.bool, device=input_ids.device).tril() |
| eye = torch.eye(T, dtype=torch.bool, device=input_ids.device) |
| attn_mask = (causal[None, None] & keep[:, None, None, :]) | eye[None, None] |
| own_gate = ~torch.cat((torch.zeros_like(keep[:, :1]), keep[:, :-1]), 1) |
| x = self.embed(input_ids) |
| for block in self.blocks: |
| x = block(x, cos, sin, attn_mask, own_gate) |
| logits = self.lm_head(self.ln_f(x)).float() |
| loss = None |
| if labels is not None: |
| loss = F.cross_entropy(logits[:, :-1].reshape(-1, logits.shape[-1]), labels[:, 1:].reshape(-1), ignore_index=-100) |
| return CausalLMOutputWithPast(loss=loss, logits=logits) |
|
|
