| 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) |
|
|
