| """transformers remote code for the exp_train hybrid LM checkpoints (trust_remote_code=True). |
| |
| tok = AutoTokenizer.from_pretrained(path) |
| model = AutoModelForCausalLM.from_pretrained(path, trust_remote_code=True, dtype=torch.bfloat16, device_map="cuda") |
| out = model.generate(**tok("The capital of France is", return_tensors="pt").to("cuda"), max_new_tokens=5) |
| |
| Architecture (exp_train / a100stack): 32 layers, 24 Gated DeltaNet + 8 gated GQA softmax attention (every 4th), |
| SwiGLU MLPs, gated residuals, Delta Block attention residuals ("repo" variant) and a hashed bigram/trigram memory |
| (ngram.safetensors, canonical-token hashing, +-1 sign trick) injected before blocks 2 and 16. |
| |
| Parameter names are the training checkpoint's, so model.safetensors loads unchanged; the n-gram table is read from |
| ngram.safetensors after loading. Matrices run in BF16, norms / gates / the residual stream in FP32 as in training. |
| Inference only, batch size 1. Prefill uses FLA's chunked gated-delta-rule kernel; decoding uses the fused recurrent |
| kernel, a preallocated KV cache, and a CUDA graph for the single-token step. |
| """ |
| import hashlib |
| import json |
| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule |
| from transformers import GenerationMixin, PretrainedConfig, PreTrainedModel |
| from transformers.modeling_outputs import CausalLMOutputWithPast |
|
|
|
|
| class HybridLMConfig(PretrainedConfig): |
| model_type = "exp_train_hybrid" |
|
|
| def __init__(self, model_config=None, modifications=None, kv_cache_len=4096, use_cuda_graph=True, |
| bos_token_id=2, eos_token_id=3, pad_token_id=0, tie_word_embeddings=False, **kwargs): |
| self.model_config = model_config or {} |
| self.modifications = modifications or {} |
| self.kv_cache_len = kv_cache_len |
| self.use_cuda_graph = use_cuda_graph |
| mc = self.model_config |
| if mc: |
| self.vocab_size = mc["vocab_size"] |
| self.hidden_size = mc["d_model"] |
| self.num_hidden_layers = mc["n_layers"] |
| self.num_attention_heads = mc["n_heads"] |
| self.num_key_value_heads = mc["n_kv_heads"] |
| self.head_dim = mc["head_dim"] |
| self.intermediate_size = mc["ffn_dim"] |
| super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, pad_token_id=pad_token_id, |
| tie_word_embeddings=tie_word_embeddings, **kwargs) |
|
|
|
|
| |
| _MULTS = [0x9E3779B97F4A7C15, 0xC2B2AE3D27D4EB4F, 0x165667B19E3779F9, 0xD6E8FEB86659FD93, |
| 0xFF51AFD7ED558CCD, 0xC4CEB9FE1A85EC53, 0x94D049BB133111EB, 0xBF58476D1CE4E5B9] |
| _SIGN_MULTS = [0xA0761D6478BD642F, 0xE7037ED1A0B428DB, 0x8EBC6AF09C88C6E3, 0x589965CC75374CC3, |
| 0x1D8E4E27C47D124F, 0x2D358DCCAA6C78A5, 0x9FB21C651E98DF25, 0xD3833E804F4C574B] |
|
|
|
|
| def _i64(x): |
| return x - (1 << 64) if x >= (1 << 63) else x |
|
|
|
|
| def _shift(ids, j): |
| return torch.cat([ids.new_zeros(j), ids[: ids.numel() - j]]) if j else ids |
|
|
|
|
| def build_canon_map(tokenizer_file): |
| """ngram_canon.py: token ids whose text normalizes to the same string share n-gram rows (hashing only).""" |
| from tokenizers import Regex, Tokenizer, normalizers |
| tokenizer = Tokenizer.from_file(str(tokenizer_file)) |
| norm = normalizers.Sequence([ |
| normalizers.NFKC(), normalizers.NFD(), normalizers.StripAccents(), normalizers.Lowercase(), |
| normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), normalizers.Replace(Regex(r"^ $"), ""), |
| normalizers.Strip(), normalizers.Replace("", " ")]) |
| vocab = tokenizer.get_vocab_size(with_added_tokens=True) |
| special = {t["id"] for t in json.loads(tokenizer.to_str()).get("added_tokens", []) if t.get("special")} |
| first, canon = {}, torch.arange(vocab, dtype=torch.int64) |
| for t in range(vocab): |
| if t in special: |
| continue |
| text = tokenizer.decode([t], skip_special_tokens=False) |
| if "�" in text: |
| continue |
| key = norm.normalize_str(text) |
| if key: |
| canon[t] = first.setdefault(key, t) |
| return canon |
|
|
|
|
| def _rms(x, w, eps): |
| return F.rms_norm(x.float(), (x.shape[-1],), w.float(), eps) |
|
|
|
|
| def _select(sources, qw, eps): |
| """Delta Block: sum_i softmax_i(q . RMSNorm_w(s_i)) s_i (qw = query * weight).""" |
| logits = torch.stack([(s * qw).sum(-1) * torch.rsqrt(s.pow(2).mean(-1) + eps) for s in sources]) |
| p = logits.softmax(0) |
| out = p[0].unsqueeze(-1) * sources[0] |
| for pi, s in zip(p[1:], sources[1:]): |
| out = out + pi.unsqueeze(-1) * s |
| return out |
|
|
|
|
| |
| def _linear(i, o): |
| return nn.Linear(i, o, bias=False) |
|
|
|
|
| class _Weight(nn.Module): |
| def __init__(self, *shape): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(*shape)) |
|
|
|
|
| class _Gate(nn.Module): |
| def __init__(self, d): |
| super().__init__() |
| self.a = nn.Parameter(torch.zeros(d)) |
|
|
|
|
| class GatedAttention(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| H, KV, hd, d = c["n_heads"], c["n_kv_heads"], c["head_dim"], c["d_model"] |
| self.q_proj, self.gate_proj, self.o_proj = _linear(d, H * hd), _linear(d, H * hd), _linear(H * hd, d) |
| self.k_proj, self.v_proj = _linear(d, KV * hd), _linear(d, KV * hd) |
| self.q_norm, self.k_norm = _Weight(hd), _Weight(hd) |
|
|
|
|
| class GatedDeltaNet(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| H, dk, d = c["gdn_heads"], c["gdn_head_dim"], c["d_model"] |
| dv, K = int(dk * c["gdn_expand_v"]), c["gdn_conv_size"] |
| self.q_proj, self.k_proj = _linear(d, H * dk), _linear(d, H * dk) |
| self.v_proj, self.g_proj, self.o_proj = _linear(d, H * dv), _linear(d, H * dv), _linear(H * dv, d) |
| self.a_proj, self.b_proj = _linear(d, H), _linear(d, H) |
| self.A_log, self.dt_bias = nn.Parameter(torch.zeros(H)), nn.Parameter(torch.zeros(H)) |
| self.q_conv1d, self.k_conv1d, self.v_conv1d = _Weight(H * dk, 1, K), _Weight(H * dk, 1, K), _Weight(H * dv, 1, K) |
| self.o_norm = _Weight(dv) |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, d, f): |
| super().__init__() |
| self.w_gate, self.w_up, self.w_down = _linear(d, f), _linear(d, f), _linear(f, d) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, c, i, has_attn_sources): |
| super().__init__() |
| self.is_attn = (i + 1) % c["full_attn_every"] == 0 |
| d = c["d_model"] |
| self.norm1, self.norm2 = _Weight(d), _Weight(d) |
| self.mixer = GatedAttention(c) if self.is_attn else GatedDeltaNet(c) |
| self.mlp = SwiGLU(d, c["ffn_dim"]) |
| self.res1, self.res2 = _Gate(d), _Gate(d) |
| self.delta_mlp_query, self.delta_mlp_weight = nn.Parameter(torch.zeros(d)), nn.Parameter(torch.ones(d)) |
| if has_attn_sources: |
| self.delta_attn_query, self.delta_attn_weight = nn.Parameter(torch.zeros(d)), nn.Parameter(torch.ones(d)) |
|
|
|
|
| class NgramMemory(nn.Module): |
| def __init__(self, ng, d_model): |
| super().__init__() |
| T, D, n = len(ng["orders"]) * ng["heads_per_order"], ng["dim"], len(ng["inject_layers"]) |
| self.proj = nn.ModuleList([_linear(T * D, d_model) for _ in range(n)]) |
| self.key_norm = nn.ModuleList([_Weight(d_model) for _ in range(n)]) |
| self.val_norm = nn.ModuleList([_Weight(d_model) for _ in range(n)]) |
| self.register_buffer("sign_pool", torch.zeros(ng["sign_pool"], D), persistent=True) |
| |
|
|
|
|
|
|
| class _State: |
| """Static decode buffers (one active sequence). Fixed addresses so the decode step can be a CUDA graph.""" |
|
|
| def __init__(self, m, max_len, device): |
| self.max_len, self.device = max_len, device |
| C = 2 * m.GH * m.GK + m.GH * m.GV |
| self.conv = [None if a else torch.zeros(C, m.CONV - 1, device=device) for a in m.is_attn] |
| self.S = [None if a else torch.zeros(1, m.GH, m.GK, m.GV, device=device) for a in m.is_attn] |
| self.k = [torch.zeros(1, m.KV, max_len, m.HD, device=device, dtype=m.mm_dtype) if a else None for a in m.is_attn] |
| self.v = [torch.zeros(1, m.KV, max_len, m.HD, device=device, dtype=m.mm_dtype) if a else None for a in m.is_attn] |
| self.ctx = torch.zeros(3, dtype=torch.int64, device=device) |
| self.pos = torch.zeros(1, dtype=torch.int64, device=device) |
| self.tok = torch.zeros(1, dtype=torch.int64, device=device) |
| self.graph = self.logits = None |
| self.owner = None |
|
|
| def recurrent(self): |
| return [t for t in self.conv + self.S + [self.ctx, self.pos] if t is not None] |
|
|
|
|
| class HybridCache: |
| """What generate() passes back as past_key_values. A new prefill takes over the model's static buffers.""" |
| is_compileable = False |
|
|
| def __init__(self, state): |
| self.state, self.seen = state, 0 |
| state.owner = self |
|
|
| def get_seq_length(self, layer_idx=0): |
| return self.seen |
|
|
|
|
| |
| class HybridLMForCausalLM(PreTrainedModel, GenerationMixin): |
| config_class = HybridLMConfig |
| base_model_prefix = "hybrid" |
| _no_split_modules = ["Block"] |
| _is_stateful = True |
| _supports_sdpa = True |
| |
| _keep_in_fp32_modules_strict = ["norm", "res1", "res2", "delta_", "A_log", "dt_bias", "embed", "sign_pool", |
| "conv1d"] |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| c, mods = config.model_config, config.modifications |
| assert mods.get("delta_block") and mods.get("delta_variant", "repo") == "repo" |
| assert mods.get("depth_attn", "off") == "off" |
| self.D, self.EPS = c["d_model"], c["norm_eps"] |
| self.H, self.KV, self.HD = c["n_heads"], c["n_kv_heads"], c["head_dim"] |
| self.GH, self.GK = c["gdn_heads"], c["gdn_head_dim"] |
| self.GV, self.CONV = int(self.GK * c["gdn_expand_v"]), c["gdn_conv_size"] |
| n = c["n_layers"] |
| self.per_block = math.ceil(n / mods.get("delta_blocks", 4)) |
| self.is_attn = [(i + 1) % c["full_attn_every"] == 0 for i in range(n)] |
| ng = c["ngram"] |
| self.ng = ng |
| self.rows = max(1024, ng["total_params"] // (len(ng["orders"]) * ng["heads_per_order"] * ng["dim"])) |
| self.inject = {l: k for k, l in enumerate(ng["inject_layers"])} |
| self.embed = nn.Embedding(c["vocab_size"], c["d_model"]) |
| |
| self.blocks = nn.ModuleList([Block(c, i, i > 0) for i in range(n)]) |
| self.norm = _Weight(c["d_model"]) |
| self.lm_head = _linear(c["d_model"], c["vocab_size"]) |
| self.ngram = NgramMemory(ng, c["d_model"]) |
| self._fused_ptr, self._state = None, None |
| self.post_init() |
|
|
| def _init_weights(self, module): |
| pass |
|
|
| def get_input_embeddings(self): |
| return self.embed |
|
|
| def get_output_embeddings(self): |
| return self.lm_head |
|
|
| @classmethod |
| def _supports_default_dynamic_cache(cls): |
| return False |
|
|
| |
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): |
| model = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) |
| from safetensors import safe_open |
| from transformers.utils import cached_file |
| hub = {k: kwargs[k] for k in ("revision", "token", "cache_dir", "local_files_only") if k in kwargs} |
| device = model.embed.weight.device |
| with safe_open(cached_file(pretrained_model_name_or_path, "ngram.safetensors", **hub), "pt", device="cpu") as f: |
| table = f.get_tensor("table") |
| expected = (len(model.ng["orders"]) * model.ng["heads_per_order"] * model.rows, model.ng["dim"]) |
| assert tuple(table.shape) == expected, (table.shape, expected) |
| canon = build_canon_map(cached_file(pretrained_model_name_or_path, "tokenizer.json", **hub)) |
| sha = model.config.modifications.get("ngram_canonical_sha256") |
| if model.ng.get("canonical_tokens"): |
| assert sha is None or hashlib.sha256(canon.numpy().tobytes()).hexdigest() == sha, \ |
| "canonical-token map differs from the one used in training" |
| else: |
| canon = torch.arange(model.config.vocab_size) |
| model.ngram.register_buffer("table", table.to(device), persistent=False) |
| model.ngram.register_buffer("canon", canon.to(device), persistent=False) |
| model.eval() |
| return model |
|
|
| @property |
| def mm_dtype(self): |
| return self.lm_head.weight.dtype |
|
|
| def _ready(self): |
| """Fuse projections into single GEMMs (the original parameters become views into the fused weight).""" |
| if self._fused_ptr == self.blocks[0].mixer.q_proj.weight.data_ptr(): |
| return |
| with torch.no_grad(): |
| for blk in self.blocks: |
| m = blk.mixer |
| names = (("q_proj", "k_proj", "v_proj", "gate_proj") if blk.is_attn else |
| ("q_proj", "k_proj", "v_proj", "a_proj", "b_proj", "g_proj")) |
| blk._in_w = self._fuse([getattr(m, n) for n in names]) |
| blk._mlp_w = self._fuse([blk.mlp.w_gate, blk.mlp.w_up]) |
| if not blk.is_attn: |
| blk._conv_w = torch.cat([m.q_conv1d.weight[:, 0], m.k_conv1d.weight[:, 0], |
| m.v_conv1d.weight[:, 0]]).float() |
| blk._neg_A = -m.A_log.float().exp() |
| blk._g1, blk._g2 = 2 * torch.sigmoid(blk.res1.a.float()), 2 * torch.sigmoid(blk.res2.a.float()) |
| blk._dmq = (blk.delta_mlp_query * blk.delta_mlp_weight).float() |
| if hasattr(blk, "delta_attn_query"): |
| blk._daq = (blk.delta_attn_query * blk.delta_attn_weight).float() |
| self._fused_ptr = self.blocks[0].mixer.q_proj.weight.data_ptr() |
| self._state = None |
|
|
| @staticmethod |
| def _fuse(linears): |
| w = torch.cat([l.weight.data for l in linears]) |
| o = 0 |
| for l in linears: |
| n = l.weight.shape[0] |
| l.weight.data = w[o:o + n] |
| o += n |
| return w |
|
|
| def _get_state(self): |
| dev = self.embed.weight.device |
| if self._state is None or self._state.device != dev: |
| self._state = _State(self, self.config.kv_cache_len, dev) |
| return self._state |
|
|
| |
| def _ngram_memory(self, cids, pos): |
| ng, rows, signs, valid = self.ng, [], [], [] |
| t = 0 |
| for order in ng["orders"]: |
| for head in range(ng["heads_per_order"]): |
| h = torch.zeros_like(cids) |
| m = _i64(_MULTS[(head * 3 + order) % 8]) |
| for j in range(order): |
| h = h * m + (_shift(cids, j) + 1 + j * 7919) |
| h = h ^ (h >> 29); h = h * _i64(0xBF58476D1CE4E5B9); h = h ^ (h >> 32) |
| rows.append(torch.remainder(h, self.rows) + t * self.rows) |
| if ng.get("sign_pool"): |
| s = torch.full_like(cids, _i64(0x632BE59BD9B4E019)) |
| m = _i64(_SIGN_MULTS[(head * 3 + order) % 8]) |
| for j in range(order): |
| s = s * m + (_shift(cids, j) + 1 + j * 104729) |
| s = s ^ (s >> 31); s = s * _i64(0x94D049BB133111EB); s = s ^ (s >> 29) |
| signs.append(s & (ng["sign_pool"] - 1)) |
| valid.append(pos >= order - 1) |
| t += 1 |
| emb = self.ngram.table[torch.stack(rows, 1)].float() |
| mult = torch.stack(valid, 1).unsqueeze(-1).float() |
| if signs: |
| mult = self.ngram.sign_pool[torch.stack(signs, 1)] * mult |
| return (emb * mult).flatten(1) |
|
|
| def _rope(self, x, pos): |
| d = x.shape[-1] |
| inv = 1.0 / (self.config.model_config["rope_theta"] ** (torch.arange(0, d, 2, device=x.device, dtype=torch.float32) / d)) |
| f = pos.float()[:, None] * inv[None] |
| cos, sin = f.cos()[:, None], f.sin()[:, None] |
| x1, x2 = x[..., : d // 2], x[..., d // 2:] |
| return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], -1) |
|
|
| def _gdn(self, blk, i, xn, st, decode): |
| m, N, H, dk, dv = blk.mixer, xn.shape[0], self.GH, self.GK, self.GV |
| qkv, a, b, gate = (xn @ blk._in_w.T).split((2 * H * dk + H * dv, H, H, H * dv), -1) |
| w = blk._conv_w |
| if decode: |
| win = torch.cat([st.conv[i], qkv.T.float()], 1) |
| y = (win * w).sum(-1, keepdim=True) |
| st.conv[i].copy_(win[:, 1:]) |
| else: |
| win = F.pad(qkv.T.float(), (self.CONV - 1, 0)) |
| y = F.conv1d(win[None], w[:, None], groups=w.shape[0])[0] |
| if st is not None: |
| st.conv[i].copy_(win[:, -(self.CONV - 1):]) |
| q, k, v = F.silu(y).T.to(xn.dtype).split((H * dk, H * dk, H * dv), -1) |
| g = (blk._neg_A * F.softplus(a.float() + m.dt_bias.float()))[None] |
| beta = b.float().sigmoid()[None] |
| kern = fused_recurrent_gated_delta_rule if decode else chunk_gated_delta_rule |
| o, S = kern(q=q.reshape(1, N, H, dk), k=k.reshape(1, N, H, dk), v=v.reshape(1, N, H, dv), g=g, beta=beta, |
| scale=dk ** -0.5, initial_state=st.S[i] if decode else None, |
| output_final_state=st is not None, use_qk_l2norm_in_kernel=True) |
| if st is not None: |
| st.S[i].copy_(S) |
| o = _rms(o[0], m.o_norm.weight, self.EPS) * F.silu(gate.float().view(N, H, dv)) |
| return o.reshape(N, -1).to(xn.dtype) @ m.o_proj.weight.T |
|
|
| def _attention(self, blk, i, xn, pos, st, decode): |
| m, N, H, KV, HD = blk.mixer, xn.shape[0], self.H, self.KV, self.HD |
| q, k, v, gate = (xn @ blk._in_w.T).split((H * HD, KV * HD, KV * HD, H * HD), -1) |
| q = self._rope(_rms(q.view(N, H, HD), m.q_norm.weight, self.EPS), pos).to(xn.dtype) |
| k = self._rope(_rms(k.view(N, KV, HD), m.k_norm.weight, self.EPS), pos).to(xn.dtype) |
| v = v.view(N, KV, HD) |
| if decode: |
| st.k[i].index_copy_(2, pos, k.transpose(0, 1)[None]) |
| st.v[i].index_copy_(2, pos, v.transpose(0, 1)[None]) |
| mask = (torch.arange(st.max_len, device=xn.device) <= pos)[None, None, None] |
| o = F.scaled_dot_product_attention(q.transpose(0, 1)[None], st.k[i], st.v[i], attn_mask=mask, |
| enable_gqa=True) |
| else: |
| if st is not None: |
| st.k[i][:, :, :N] = k.transpose(0, 1)[None] |
| st.v[i][:, :, :N] = v.transpose(0, 1)[None] |
| o = F.scaled_dot_product_attention(q.transpose(0, 1)[None], k.transpose(0, 1)[None], |
| v.transpose(0, 1)[None], is_causal=True, enable_gqa=True) |
| o = o[0].transpose(0, 1) * torch.sigmoid(gate.view(N, H, HD)) |
| return o.reshape(N, -1) @ m.o_proj.weight.T |
|
|
| def _run(self, ids, st, decode, keep): |
| """ids [N] (prefill from position 0) or [1] (decode at st.pos). Returns logits [keep or N, V] (FP32).""" |
| N, dt, eps = ids.shape[0], self.mm_dtype, self.EPS |
| cids = self.ngram.canon[ids] |
| if decode: |
| pos = st.pos |
| st.ctx.copy_(torch.cat([st.ctx[1:], cids])) |
| mem = self._ngram_memory(st.ctx, torch.cat([pos - 2, pos - 1, pos]))[-1:] |
| else: |
| pos = torch.arange(N, device=ids.device) |
| mem = self._ngram_memory(cids, pos) |
| if st is not None: |
| st.ctx.copy_(F.pad(cids, (3, 0))[-3:]) |
| st.pos.fill_(N) |
| x = self.embed.weight[ids].float() |
| sources = [] |
| for i, blk in enumerate(self.blocks): |
| if i in self.inject: |
| k = self.inject[i] |
| e = (mem.to(dt) @ self.ngram.proj[k].weight.T).float() |
| gate = torch.sigmoid((_rms(x, self.ngram.key_norm[k].weight, eps) * |
| _rms(e, self.ngram.val_norm[k].weight, eps)).sum(-1, keepdim=True) / math.sqrt(self.D)) |
| x = x + gate * e |
| attn_sources = list(sources) |
| if i % self.per_block == 0: |
| sources.append(x if not sources else x - sources[-1]) |
| h = x + _select(attn_sources, blk._daq, eps) if attn_sources else x |
| xn = _rms(h, blk.norm1.weight, eps).to(dt) |
| y = self._attention(blk, i, xn, pos, st, decode) if blk.is_attn else self._gdn(blk, i, xn, st, decode) |
| x = x + blk._g1 * y.float() |
| h = x + _select(sources, blk._dmq, eps) |
| gt, up = (_rms(h, blk.norm2.weight, eps).to(dt) @ blk._mlp_w.T).chunk(2, -1) |
| x = x + blk._g2 * ((F.silu(gt) * up) @ blk.mlp.w_down.weight.T).float() |
| if keep: |
| x = x[-keep:] |
| logits = (_rms(x, self.norm.weight, eps).to(dt) @ self.lm_head.weight.T).float() |
| if decode: |
| st.pos += 1 |
| return logits |
|
|
| def _decode(self, st, token): |
| st.tok.copy_(token) |
| if not (self.config.use_cuda_graph and st.device.type == "cuda"): |
| return self._run(st.tok, st, True, 0) |
| if st.graph is None: |
| snap = [t.clone() for t in st.recurrent()] |
| s = torch.cuda.Stream() |
| s.wait_stream(torch.cuda.current_stream()) |
| with torch.cuda.stream(s): |
| for _ in range(2): |
| self._run(st.tok, st, True, 0) |
| torch.cuda.current_stream().wait_stream(s) |
| graph = torch.cuda.CUDAGraph() |
| with torch.cuda.graph(graph): |
| st.logits = self._run(st.tok, st, True, 0) |
| for dst, src in zip(st.recurrent(), snap): |
| dst.copy_(src) |
| st.graph = graph |
| st.graph.replay() |
| return st.logits.clone() |
|
|
| |
| def forward(self, input_ids=None, attention_mask=None, past_key_values=None, use_cache=None, labels=None, |
| logits_to_keep=0, return_dict=True, **kwargs): |
| if input_ids is None or input_ids.shape[0] != 1: |
| raise ValueError("this model implementation supports input_ids of batch size 1") |
| if attention_mask is not None and not bool(attention_mask.all()): |
| raise ValueError("padding is not supported (batch size 1)") |
| self._ready() |
| use_cache = self.config.use_cache if use_cache is None else use_cache |
| ids = input_ids[0] |
| keep = logits_to_keep if isinstance(logits_to_keep, int) else 0 |
| cache = past_key_values if isinstance(past_key_values, HybridCache) else None |
| with torch.no_grad(): |
| if cache is not None and cache.seen > 0: |
| assert cache.state.owner is cache, "this cache was superseded by a newer prefill" |
| if cache.seen + ids.numel() > cache.state.max_len: |
| raise ValueError(f"sequence exceeds kv_cache_len={cache.state.max_len}") |
| logits = torch.cat([self._decode(cache.state, t) for t in ids.split(1)]) |
| cache.seen += ids.numel() |
| else: |
| cache = None |
| if use_cache: |
| if ids.numel() > self.config.kv_cache_len: |
| raise ValueError(f"prompt exceeds kv_cache_len={self.config.kv_cache_len}") |
| cache = HybridCache(self._get_state()) |
| cache.seen = ids.numel() |
| logits = self._run(ids, cache.state if cache else None, False, 0 if labels is not None else keep) |
| loss = None |
| if labels is not None: |
| loss = F.cross_entropy(logits[:-1], labels[0, 1:], ignore_index=-100) |
| return CausalLMOutputWithPast(loss=loss, logits=logits[None], past_key_values=cache) |
|
|
| def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, use_cache=True, |
| **kwargs): |
| if isinstance(past_key_values, HybridCache) and past_key_values.seen > 0: |
| input_ids = input_ids[:, past_key_values.seen:] |
| return dict(input_ids=input_ids, past_key_values=past_key_values, use_cache=use_cache, logits_to_keep=1) |
|
|
|
|
| HybridLMConfig.register_for_auto_class() |
| HybridLMForCausalLM.register_for_auto_class("AutoModelForCausalLM") |
|
|
