File size: 9,960 Bytes
4c0a0c6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | import math
import torch
import torch.nn.functional as F
from einops import rearrange
from torch import nn
from transformers import PreTrainedModel, PretrainedConfig
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
try:
from fla.layers.utils import get_layer_cache, update_layer_cache
from fla.models.utils import Cache, FLAUnsupportedCacheGenerationMixin
from fla.modules import FusedRMSNormSwishGate, RMSNorm, ShortConvolution
from fla.ops.gdn2 import chunk_gdn2, fused_recurrent_gdn2
except (ImportError, AttributeError) as e:
raise ImportError('Linear-50M requires a current CUDA build of flash-linear-attention: pip install "flash-linear-attention[cuda]"') from e
class Linear50MConfig(PretrainedConfig):
model_type = "linear_50m"
def __init__(self, vocab_size=32000, hidden_size=384, num_hidden_layers=16, num_heads=6,
head_dim=64, intermediate_size=1120, conv_size=4, norm_eps=1e-5,
max_position_embeddings=4096, use_cache=True, **kwargs):
for key, value in dict(bos_token_id=1, eos_token_id=2, pad_token_id=2,
tie_word_embeddings=True).items(): kwargs.setdefault(key, value)
super().__init__(**kwargs); self.vocab_size, self.hidden_size, self.num_hidden_layers = vocab_size, hidden_size, num_hidden_layers
self.num_heads, self.head_dim, self.intermediate_size = num_heads, head_dim, intermediate_size
self.conv_size, self.norm_eps, self.max_position_embeddings, self.use_cache = conv_size, norm_eps, max_position_embeddings, use_cache
class NotGDN2(nn.Module):
def __init__(self, c, layer_idx):
super().__init__(); d, h, hd = c.hidden_size, c.num_heads, c.head_dim; self.h, self.hd, self.layer_idx = h, hd, layer_idx
self.qkv = nn.Linear(d, 3*d, bias=False); self.conv = ShortConvolution(3*d, c.conv_size, bias=False, activation="silu")
self.f = nn.Sequential(nn.Linear(d, hd, False), nn.Linear(hd, d, False)); self.bw = nn.Linear(d, 2*d, bias=False)
self.gate = nn.Sequential(nn.Linear(d, hd, False), nn.Linear(hd, d, True)); self.recover = nn.Sequential(nn.Linear(d, hd, False), nn.Linear(hd, d, True))
self.A_log = nn.Parameter(torch.empty(h).uniform_(1, 16).log()); dt = torch.exp(torch.empty(d).uniform_(math.log(.001), math.log(.1)))
self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt))); self.A_log._no_weight_decay = self.dt_bias._no_weight_decay = True
self.norm = FusedRMSNormSwishGate(hd, eps=c.norm_eps); self.out = nn.Linear(d, d, bias=False)
self.apply(self._init)
for weight, parts in ((self.qkv.weight, 3), (self.bw.weight, 2)):
for part in weight.chunk(parts, 0): nn.init.xavier_uniform_(part, gain=2**-2.5)
nn.init.zeros_(self.recover[-1].weight); nn.init.zeros_(self.recover[-1].bias)
@staticmethod
def _init(m):
if getattr(m, "_is_hf_initialized", False): return
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight, gain=2**-2.5)
if m.bias is not None: nn.init.zeros_(m.bias)
m._is_hf_initialized = True
def forward(self, x, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
length = x.shape[1]; last = get_layer_cache(self, past_key_values); conv = None if last is None else last["conv_state"]; cu = kwargs.get("cu_seqlens")
qkv, conv = self.conv(self.qkv(x), cache=conv, output_final_state=use_cache, cu_seqlens=cu)
q, k, v = (rearrange(z, "b t (h d) -> b t h d", h=self.h) for z in qkv.chunk(3, -1))
g = rearrange(-self.A_log.float().exp().repeat_interleave(self.hd) *
F.softplus(self.f(x).float() + self.dt_bias.float()),
"b t (h d) -> b t h d", h=self.h)
b, w = (rearrange(z.sigmoid(), "b t (h d) -> b t h d", h=self.h) for z in self.bw(x).chunk(2, -1))
state = None if last is None else last["recurrent_state"]
op = fused_recurrent_gdn2 if use_cache and not self.training and not torch.is_grad_enabled() and x.shape[1] <= 64 else chunk_gdn2
both, state = op(*(torch.cat(z, 2) for z in ((q, q), (k, k), (v, v),
(g, torch.zeros_like(g)), (b, torch.zeros_like(b)), (w, w))),
scale=self.hd**-.5, initial_state=state, output_final_state=use_cache,
use_qk_l2norm_in_kernel=True, use_gate_in_kernel=False, cu_seqlens=cu)
a, archive = both.chunk(2, 2); protected = archive-a
support = protected*((a.float().square().mean(-1, True)+1e-6)/(protected.float().square().mean(-1, True)+1e-6)).sqrt().to(protected)
theta = rearrange(self.recover(x), "b t (h d) -> b t h d", h=self.h); o = a + theta.sin()*support
if use_cache: update_layer_cache(self, past_key_values, recurrent_state=state, conv_state=conv, offset=length)
z = rearrange(self.gate(x), "b t (h d) -> b t h d", h=self.h); return self.out(rearrange(self.norm(o, z), "b t h d -> b t (h d)"))
class SwiGLU(nn.Module):
def __init__(self, c): super().__init__(); self.up, self.down = nn.Linear(c.hidden_size, 2*c.intermediate_size, False), nn.Linear(c.intermediate_size, c.hidden_size, False)
def forward(self, x): a, b = self.up(x).chunk(2, -1); return self.down(F.silu(a)*b)
class Block(nn.Module):
def __init__(self, c, i): super().__init__(); self.n1, self.mix, self.n2, self.mlp = RMSNorm(c.hidden_size), NotGDN2(c, i), RMSNorm(c.hidden_size), SwiGLU(c)
def forward(self, x, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
x = x + self.mix(self.n1(x), attention_mask, past_key_values, use_cache, **kwargs); return x + self.mlp(self.n2(x))
class Linear50MPreTrainedModel(PreTrainedModel):
config_class, base_model_prefix, _supports_cache_class, _no_split_modules = Linear50MConfig, "model", True, ["Block"]
def _init_weights(self, m):
if getattr(m, "_is_hf_initialized", False): return
if isinstance(m, (nn.Linear, nn.Embedding)):
nn.init.normal_(m.weight, std=.02)
if getattr(m, "bias", None) is not None: nn.init.zeros_(m.bias)
class Linear50MModel(Linear50MPreTrainedModel):
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.norm = RMSNorm(config.hidden_size); self.post_init()
def get_input_embeddings(self): return self.embed
def set_input_embeddings(self, value): self.embed = value
def forward(self, input_ids=None, attention_mask=None, inputs_embeds=None, past_key_values=None, use_cache=None,
output_hidden_states=None, return_dict=None, **kwargs):
if (input_ids is None) == (inputs_embeds is None): raise ValueError("Pass exactly one of input_ids or inputs_embeds")
x = self.embed(input_ids) if inputs_embeds is None else inputs_embeds
if x.device.type != "cuda": raise RuntimeError("Linear-50M currently requires CUDA/FLA inference")
if attention_mask is not None:
if past_key_values is None and not bool(attention_mask.all()): raise ValueError("Variable-length padded batches are not supported yet")
attention_mask = None
output_hidden_states = self.config.output_hidden_states if output_hidden_states is None else output_hidden_states
return_dict = self.config.return_dict if return_dict is None else return_dict
use_cache = getattr(self.config, "use_cache", True) if use_cache is None and not self.training else bool(use_cache)
if use_cache and not isinstance(past_key_values, Cache): past_key_values = Cache.from_legacy_cache(past_key_values)
hidden = () if output_hidden_states else None
for block in self.blocks:
if output_hidden_states: hidden += (x,)
x = block(x, attention_mask, past_key_values, use_cache, **kwargs)
x = self.norm(x)
if output_hidden_states: hidden += (x,)
if not return_dict: return tuple(v for v in (x, past_key_values, hidden) if v is not None)
return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=past_key_values, hidden_states=hidden)
class Linear50MForCausalLM(Linear50MPreTrainedModel, FLAUnsupportedCacheGenerationMixin):
_tied_weights_keys = {"lm_head.weight": "model.embed.weight"}
def __init__(self, config): super().__init__(config); self.model = Linear50MModel(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
def set_input_embeddings(self, value): self.model.embed = value
def get_output_embeddings(self): return self.lm_head
def set_output_embeddings(self, value): self.lm_head = value
def get_decoder(self): return self.model
def set_decoder(self, value): self.model = value
def forward(self, input_ids=None, attention_mask=None, inputs_embeds=None, past_key_values=None, labels=None,
use_cache=None, output_hidden_states=None, return_dict=None, logits_to_keep=0, **kwargs):
return_dict = self.config.return_dict if return_dict is None else return_dict
out = self.model(input_ids, attention_mask, inputs_embeds, past_key_values, use_cache, output_hidden_states, return_dict, **kwargs); h = out[0]
logits = self.lm_head(h if labels is not None or not logits_to_keep else h[:, -logits_to_keep:]); loss = None
if labels is not None: loss = F.cross_entropy(logits[:, :-1].float().flatten(0, 1), labels[:, 1:].flatten(), ignore_index=-100)
if not return_dict: return ((loss,) if loss is not None else ()) + (logits,) + out[1:]
return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=out.past_key_values, hidden_states=out.hidden_states)
|