| """Penny (DeepSeek-V4.1-Flash port) for Hugging Face transformers, via trust_remote_code. |
| |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| tok = AutoTokenizer.from_pretrained("LocalWisdom/Penny-Alpha") |
| model = AutoModelForCausalLM.from_pretrained("LocalWisdom/Penny-Alpha", trust_remote_code=True, dtype="float32") |
| ids = tok(tok.eos_token + "The history of the bicycle begins", add_special_tokens=False, return_tensors="pt") |
| out = model.generate(**ids, max_new_tokens=200, do_sample=True, temperature=0.8, top_k=50, use_cache=False) |
| print(tok.decode(out[0])) |
| |
| Inference-only port of the training model (no KV cache: every step re-runs the whole sequence, which is |
| fine at this size). Weights use the official DeepSeek-V4.1 tensor names, so `model.safetensors` loads as is. |
| Penny is a base model (no chat template). Prompts must start with EOS (marks a fresh document); the tokenizer does not add it, so prepend tok.eos_token. |
| """ |
|
|
| import math |
| from dataclasses import dataclass |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
| from transformers import AutoTokenizer, GenerationMixin, PretrainedConfig, PreTrainedModel |
| from transformers.modeling_outputs import CausalLMOutput |
|
|
|
|
| |
| |
|
|
|
|
| class DeepseekV41Config(PretrainedConfig): |
| """Wraps the nested config.json: top-level HF fields + `text_config` (the architecture) + `penny`.""" |
|
|
| model_type = "deepseek_v41" |
|
|
| def __init__(self, text_config=None, penny=None, **kwargs): |
| super().__init__(**kwargs) |
| |
| self.arch = dict(text_config or {}) |
| self.penny = dict(penny or {}) |
|
|
| def to_dict(self): |
| d = super().to_dict() |
| d["text_config"] = d.pop("arch", self.arch) |
| return d |
|
|
|
|
| @dataclass |
| class Args: |
| vocab_size: int |
| dim: int |
| moe_inter_dim: int |
| n_layers: int |
| n_heads: int |
| n_routed_experts: int |
| n_shared_experts: int |
| n_activated_experts: int |
| score_func: str |
| norm_topk_prob: bool |
| route_scale: float |
| swiglu_limit: float |
| q_lora_rank: int |
| head_dim: int |
| rope_head_dim: int |
| norm_eps: float |
| o_groups: int |
| o_lora_rank: int |
| window_size: int |
| compress_ratios: tuple |
| kv_source_layers: tuple |
| index_source_layers: tuple |
| compress_rope_theta: float |
| original_seq_len: int |
| rope_theta: float |
| rope_factor: float |
| beta_fast: int |
| beta_slow: int |
| index_n_heads: int |
| index_head_dim: int |
| index_topk: int |
| candidate_source_layer: int |
| candidate_topk_blocks: int |
| candidate_block_size: int |
| hc_mult: int |
| hc_sinkhorn_iters: int |
| hc_eps: float |
| engram_layer_ids: tuple |
| engram_num_embeddings: tuple |
| engram_max_ngram_size: int |
| engram_vocab_size: int |
| engram_n_heads: int |
| engram_head_dim: int |
| engram_pad_id: int |
| engram_compressed_vocab_size: int |
| max_seq_len: int |
| attn_chunk: int = 256 |
|
|
| @classmethod |
| def from_config(cls, cfg: DeepseekV41Config) -> "Args": |
| t = cfg.arch |
| rs = t.get("rope_scaling") or {} |
| n = t["num_hidden_layers"] |
| return cls( |
| vocab_size=t["vocab_size"], dim=t["hidden_size"], moe_inter_dim=t["moe_intermediate_size"], |
| n_layers=n, n_heads=t["num_attention_heads"], n_routed_experts=t["n_routed_experts"], |
| n_shared_experts=t["n_shared_experts"], n_activated_experts=t["num_experts_per_tok"], |
| score_func=t["scoring_func"], norm_topk_prob=t["norm_topk_prob"], |
| route_scale=t["routed_scaling_factor"], swiglu_limit=t["swiglu_limit"], |
| q_lora_rank=t["q_lora_rank"], head_dim=t["head_dim"], rope_head_dim=t["qk_rope_head_dim"], |
| norm_eps=t["rms_norm_eps"], o_groups=t["o_groups"], o_lora_rank=t["o_lora_rank"], |
| window_size=t["sliding_window"], compress_ratios=tuple(t["compress_ratios"][:n]), |
| kv_source_layers=tuple(t["kv_source_layer_ids"]), index_source_layers=tuple(t["index_source_layer_ids"]), |
| compress_rope_theta=t["compress_rope_theta"], |
| original_seq_len=rs.get("original_max_position_embeddings", 0), rope_theta=t["rope_theta"], |
| rope_factor=rs.get("factor", 1), beta_fast=rs.get("beta_fast", 32), beta_slow=rs.get("beta_slow", 1), |
| index_n_heads=t["index_n_heads"], index_head_dim=t["index_head_dim"], index_topk=t["index_topk"], |
| candidate_source_layer=t["candidate_source_layer_id"], candidate_topk_blocks=t["candidate_topk_blocks"], |
| candidate_block_size=t["candidate_block_size"], hc_mult=t["hc_mult"], |
| hc_sinkhorn_iters=t["hc_sinkhorn_iters"], hc_eps=t["hc_eps"], |
| engram_layer_ids=tuple(t["engram_layer_ids"]), engram_num_embeddings=tuple(t["engram_num_embeddings"]), |
| engram_max_ngram_size=t["engram_max_ngram_size"], engram_vocab_size=t["engram_vocab_size"], |
| engram_n_heads=t["engram_n_heads"], engram_head_dim=t["engram_head_dim"], |
| engram_pad_id=t["engram_pad_token_id"], engram_compressed_vocab_size=t["engram_compressed_vocab_size"], |
| max_seq_len=cfg.penny.get("trained_seq_len", 8192), |
| ) |
|
|
|
|
| |
| |
|
|
|
|
| def _is_prime(n: int) -> bool: |
| if n < 2: |
| return False |
| i = 2 |
| while i * i <= n: |
| if n % i == 0: |
| return False |
| i += 1 |
| return True |
|
|
|
|
| def _next_prime(start: int, seen: set) -> int: |
| c = start + 1 |
| while not _is_prime(c) or c in seen: |
| c += 1 |
| return c |
|
|
|
|
| def _compressed_token_map(tokenizer): |
| """Collapse tokens that normalize alike (" The", "the", "THE") onto one compressed id.""" |
| from tokenizers import Regex, normalizers |
|
|
| sentinel = "\ue000" |
| normalizer = normalizers.Sequence([ |
| normalizers.NFKC(), normalizers.NFD(), normalizers.StripAccents(), normalizers.Lowercase(), |
| normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), normalizers.Replace(Regex(r"^ $"), sentinel), |
| normalizers.Strip(), normalizers.Replace(sentinel, " "), |
| ]) |
| backend = tokenizer.backend_tokenizer |
| key_to_new, lookup = {}, [0] * len(tokenizer) |
| for token_id in range(len(tokenizer)): |
| text = backend.decode([token_id], skip_special_tokens=False) |
| if "\ufffd" in text: |
| key = backend.id_to_token(token_id) |
| else: |
| normalized = normalizer.normalize_str(text) |
| key = normalized if normalized else text |
| new_id = key_to_new.get(key) |
| if new_id is None: |
| new_id = len(key_to_new) |
| key_to_new[key] = new_id |
| lookup[token_id] = new_id |
| return lookup, len(key_to_new) |
|
|
|
|
| class NgramHash: |
| """input_ids [B, L] -> hash ids [B, L, n_engram_layers, (max_ngram_size-1) * n_heads]. Built lazily |
| (needs the tokenizer) and not an nn.Module, so it never touches the weights.""" |
|
|
| def __init__(self, a: Args, tokenizer): |
| layer_ids = tuple(a.engram_layer_ids) |
| self.max_n = a.engram_max_ngram_size |
| primes, seen = [], set() |
| for _ in layer_ids: |
| per_ngram = [] |
| for _ in range(self.max_n - 1): |
| sizes, current = [], a.engram_vocab_size - 1 |
| for _ in range(a.engram_n_heads): |
| current = _next_prime(current, seen) |
| seen.add(current) |
| sizes.append(current) |
| per_ngram.append(sizes) |
| primes.append(per_ngram) |
| token_map, vocab_size = _compressed_token_map(tokenizer) |
| assert vocab_size == a.engram_compressed_vocab_size, (vocab_size, a.engram_compressed_vocab_size) |
| self.pad_id = token_map[a.engram_pad_id] |
| flat = [[p for per_ngram in layer for p in per_ngram] for layer in primes] |
| offsets = [np.cumsum([0, *sizes[:-1]]) for sizes in flat] |
| max_long = np.iinfo(np.int64).max |
| bound = max(1, (max_long // vocab_size) // 2) |
| mult = [] |
| for layer_id in layer_ids: |
| g = np.random.default_rng(10007 * layer_id) |
| mult.append(torch.tensor(g.integers(low=0, high=bound, size=(self.max_n,), dtype=np.int64) * 2 + 1)) |
| self.primes = torch.tensor(primes) |
| self.offsets = torch.tensor(np.array(offsets)) |
| self.multipliers = torch.stack(mult) |
| self.token_map = torch.tensor(token_map) |
|
|
| @torch.no_grad() |
| def __call__(self, input_ids: torch.Tensor) -> torch.Tensor: |
| dev = input_ids.device |
| primes, offsets, multipliers, token_map = ( |
| self.primes.to(dev), self.offsets.to(dev), self.multipliers.to(dev), self.token_map.to(dev)) |
| compressed = token_map[input_ids] |
| tokens = [compressed] |
| for shift in range(1, self.max_n): |
| shifted = torch.full_like(compressed, self.pad_id) |
| shifted[:, shift:] = compressed[:, :-shift] |
| tokens.append(shifted) |
| tokens = torch.stack(tokens, dim=-1) |
| products = tokens.unsqueeze(2) * multipliers |
| rolling, hashes = products[..., 0], [] |
| for i in range(1, self.max_n): |
| rolling = torch.bitwise_xor(rolling, products[..., i]) |
| hashes.append(rolling.unsqueeze(-1) % primes[:, i - 1]) |
| return torch.cat(hashes, dim=-1) + offsets |
|
|
|
|
| |
| |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim: int, eps: float): |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x): |
| dtype = x.dtype |
| x = x.float() |
| x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps) |
| return (self.weight.float() * x).to(dtype) |
|
|
|
|
| def precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow, device): |
| freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) |
| if original_seq_len > 0: |
| def corrected_dim(rotations): |
| return dim * math.log(original_seq_len / (rotations * 2 * math.pi)) / (2 * math.log(base)) |
|
|
| low = max(math.floor(corrected_dim(beta_fast)), 0) |
| high = min(math.ceil(corrected_dim(beta_slow)), dim - 1) |
| ramp = ((torch.arange(dim // 2, dtype=torch.float32) - low) / max(high - low, 1e-3)).clamp(0, 1) |
| smooth = 1 - ramp |
| freqs = freqs / factor * (1 - smooth) + freqs * smooth |
| freqs = torch.outer(torch.arange(seqlen), freqs) |
| return torch.polar(torch.ones_like(freqs), freqs).to(device) |
|
|
|
|
| def rope(x, freqs_cis, rd, inverse=False): |
| """Rotate the last `rd` dims of x (adjacent pairs as complex numbers). x: [b,s,d] or [b,s,h,d].""" |
| nope, pe = x[..., :-rd], x[..., -rd:] |
| z = torch.view_as_complex(pe.float().unflatten(-1, (-1, 2)).contiguous()) |
| f = freqs_cis.conj() if inverse else freqs_cis |
| f = f.view(1, z.size(1), *([1] * (z.ndim - 3)), z.size(-1)) |
| pe = torch.view_as_real(z * f).flatten(-2).to(x.dtype) |
| return torch.cat([nope, pe], dim=-1) |
|
|
|
|
| def idxs_to_mask(idxs, n): |
| """[b,m,k] int indices (-1 = empty) -> [b,m,n] bool mask.""" |
| mask = torch.zeros(*idxs.shape[:2], n + 1, dtype=torch.bool, device=idxs.device) |
| mask.scatter_(-1, torch.where(idxs < 0, n, idxs).long(), True) |
| return mask[..., :n] |
|
|
|
|
| def chunked_sink_attention(q, kv, attn_sink, window, scale, comp_kv=None, comp_idxs=None, chunk=256): |
| """Sliding-window attention over raw KV plus top-k compressed KV, with a learned sink logit. |
| q [b,L,h,d]; kv [b,L,d]; comp_kv [b,T,d]; comp_idxs [b,L,k] into comp_kv (-1 = empty).""" |
| L = q.size(1) |
| outs = [] |
| for s0 in range(0, L, chunk): |
| s1 = min(s0 + chunk, L) |
| k0 = max(0, s0 - window + 1) |
| qpos = torch.arange(s0, s1, device=q.device).unsqueeze(1) |
| kpos = torch.arange(k0, s1, device=q.device) |
| wmask = (kpos <= qpos) & (kpos > qpos - window) |
| qc, kw = q[:, s0:s1], kv[:, k0:s1] |
| s = torch.einsum("bmhd,bnd->bhmn", qc, kw).float() * scale |
| s = s.masked_fill(~wmask, -torch.inf) |
| vals = kw |
| if comp_kv is not None: |
| cmask = idxs_to_mask(comp_idxs[:, s0:s1], comp_kv.size(1)) |
| sc = torch.einsum("bmhd,bnd->bhmn", qc, comp_kv).float() * scale |
| s = torch.cat([s, sc.masked_fill(~cmask.unsqueeze(1), -torch.inf)], dim=-1) |
| vals = torch.cat([kw, comp_kv], dim=1) |
| sink = attn_sink.float().view(1, -1, 1, 1).expand(*s.shape[:3], 1) |
| lse = torch.logsumexp(torch.cat([s, sink], dim=-1), dim=-1) |
| p = torch.exp(s - lse.unsqueeze(-1)) |
| outs.append(torch.einsum("bhmn,bnd->bmhd", p.to(vals.dtype), vals)) |
| return torch.cat(outs, dim=1) |
|
|
|
|
| def hc_split_sinkhorn(mixes, hc_scale, hc_base, hc, iters, eps): |
| pre = torch.sigmoid(mixes[..., :hc] * hc_scale[0] + hc_base[:hc]) + eps |
| post = 2 * torch.sigmoid(mixes[..., hc : 2 * hc] * hc_scale[1] + hc_base[hc : 2 * hc]) |
| comb = mixes[..., 2 * hc :] * hc_scale[2] + hc_base[2 * hc :] |
| comb = comb.unflatten(-1, (hc, hc)).softmax(-1) + eps |
| comb = comb / (comb.sum(-2, keepdim=True) + eps) |
| for _ in range(iters - 1): |
| comb = comb / (comb.sum(-1, keepdim=True) + eps) |
| comb = comb / (comb.sum(-2, keepdim=True) + eps) |
| return pre, post, comb |
|
|
|
|
| |
| |
|
|
|
|
| class Shared: |
| """Per-forward hand-off between source layers and the layers that reuse their results.""" |
|
|
| def __init__(self): |
| self.compress_kv = None |
| self.index_k = None |
| self.topk_idxs = None |
| self.candidates = None |
|
|
|
|
| class Compressor(nn.Module): |
| """Pools `ratio` consecutive tokens into one KV latent with a learned softmax gate (pre-RoPE).""" |
|
|
| def __init__(self, a: Args, layer_id: int): |
| super().__init__() |
| self.ratio = a.compress_ratios[layer_id] |
| self.norm = RMSNorm(a.head_dim, a.norm_eps) |
| self.wkv = nn.Linear(a.dim, a.head_dim, bias=False) |
| self.wgate = nn.Linear(a.dim, a.head_dim, bias=False) if self.ratio > 1 else None |
|
|
| def forward(self, x): |
| if self.ratio == 1: |
| return self.norm(self.wkv(x)) |
| cutoff = x.size(1) - x.size(1) % self.ratio |
| xf = x[:, :cutoff].float() |
| kv = F.linear(xf, self.wkv.weight.float()).unflatten(1, (-1, self.ratio)) |
| score = F.linear(xf, self.wgate.weight.float()).unflatten(1, (-1, self.ratio)) |
| kv = (kv * score.softmax(dim=2)).sum(dim=2) |
| return self.norm(kv.to(x.dtype)) |
|
|
|
|
| def select_candidate_blocks(logits, compress_lens, topk_blocks, block_size): |
| width = logits.size(-1) |
| scores = F.pad(logits, (0, -width % block_size), value=-torch.inf) |
| scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1) |
| num_blocks = scores.size(-1) |
| last = (compress_lens - 1) // block_size |
| scores = scores.masked_fill(torch.arange(num_blocks, device=logits.device) == last, torch.inf) |
| top = scores.topk(min(topk_blocks, num_blocks), dim=-1) |
| keep = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, top.indices, top.values > -torch.inf) |
| return keep.repeat_interleave(block_size, dim=-1)[..., :width] |
|
|
|
|
| class Indexer(nn.Module): |
| """Scores compressed positions and keeps the `index_topk` best per query.""" |
|
|
| def __init__(self, a: Args, layer_id: int): |
| super().__init__() |
| self.owns_k = layer_id in a.kv_source_layers |
| self.ratio = a.compress_ratios[layer_id] |
| self.is_candidate_source = layer_id == a.candidate_source_layer |
| self.uses_candidates = 0 <= a.candidate_source_layer < layer_id |
| self.candidate_topk_blocks = a.candidate_topk_blocks |
| self.candidate_block_size = a.candidate_block_size |
| self.n_heads = a.index_n_heads |
| self.index_head_dim = a.index_head_dim |
| self.rd = a.rope_head_dim |
| self.index_topk = a.index_topk |
| self.softmax_scale = self.index_head_dim**-0.5 |
| self.wq_b = nn.Linear(a.q_lora_rank, self.n_heads * self.index_head_dim, bias=False) |
| self.weights_proj = nn.Linear(a.dim, self.n_heads, bias=False) |
| if self.owns_k: |
| self.wk = nn.Linear(a.head_dim, self.index_head_dim, bias=False) |
| self.k_norm = RMSNorm(self.index_head_dim, a.norm_eps) |
|
|
| @torch.no_grad() |
| def forward(self, x, qr, latent, freqs_cis, shared: Shared, chunk: int): |
| """Returns selected compressed indices [b,s,k] (-1 = empty).""" |
| seqlen, ratio = x.size(1), self.ratio |
| if self.owns_k: |
| k = self.k_norm(self.wk(latent)) |
| shared.index_k = rope(k, freqs_cis[: seqlen - seqlen % ratio : ratio], self.rd) |
| q = rope(self.wq_b(qr).unflatten(-1, (self.n_heads, self.index_head_dim)), freqs_cis[:seqlen], self.rd) |
| index_k = shared.index_k |
| weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads**-0.5) |
| T = index_k.size(1) |
| topk = min(self.index_topk, T) |
| idx_chunks, cand_chunks = [], [] |
| for s0 in range(0, seqlen, chunk): |
| s1 = min(s0 + chunk, seqlen) |
| score = torch.einsum("bshd,btd->bsht", q[:, s0:s1], index_k) |
| score = (score.relu() * weights[:, s0:s1].unsqueeze(-1)).sum(dim=2).float() |
| lens = (torch.arange(s0 + 1, s1 + 1, device=x.device) // ratio).unsqueeze(-1) |
| score = score.masked_fill(torch.arange(T, device=x.device) >= lens, -torch.inf) |
| if self.is_candidate_source: |
| cand_chunks.append(select_candidate_blocks( |
| score, lens, self.candidate_topk_blocks, self.candidate_block_size)) |
| elif self.uses_candidates and shared.candidates is not None: |
| score = score.masked_fill(~shared.candidates[:, s0:s1], -torch.inf) |
| idxs = score.topk(topk, dim=-1, sorted=False).indices |
| idx_chunks.append(torch.where(idxs < lens, idxs, -1)) |
| if self.is_candidate_source: |
| shared.candidates = torch.cat(cand_chunks, dim=1) |
| return torch.cat(idx_chunks, dim=1) |
|
|
|
|
| class Attention(nn.Module): |
| """Sliding window over raw KV, plus (compress_ratio > 0) top-k compressed KV further back. |
| Low-rank Q; grouped low-rank output projection; shared K=V latent (MQA).""" |
|
|
| def __init__(self, layer_id: int, a: Args): |
| super().__init__() |
| self.n_heads, self.head_dim, self.rd = a.n_heads, a.head_dim, a.rope_head_dim |
| self.n_groups, self.o_lora_rank, self.window = a.o_groups, a.o_lora_rank, a.window_size |
| self.ratio = a.compress_ratios[layer_id] |
| self.softmax_scale = self.head_dim**-0.5 |
| self.attn_chunk = a.attn_chunk |
| self.attn_sink = nn.Parameter(torch.zeros(self.n_heads)) |
| self.wq_a = nn.Linear(a.dim, a.q_lora_rank, bias=False) |
| self.q_norm = RMSNorm(a.q_lora_rank, a.norm_eps) |
| self.wq_b = nn.Linear(a.q_lora_rank, self.n_heads * self.head_dim, bias=False) |
| self.wkv = nn.Linear(a.dim, self.head_dim, bias=False) |
| self.kv_norm = RMSNorm(self.head_dim, a.norm_eps) |
| self.wo_a = nn.Linear(self.n_heads * self.head_dim // self.n_groups, self.n_groups * self.o_lora_rank, bias=False) |
| self.wo_b = nn.Linear(self.n_groups * self.o_lora_rank, a.dim, bias=False) |
| self.is_kv_source = layer_id in a.kv_source_layers |
| self.is_index_source = layer_id in a.index_source_layers |
| self.compressor = Compressor(a, layer_id) if self.is_kv_source else None |
| self.indexer = Indexer(a, layer_id) if self.is_index_source else None |
| |
| if self.ratio: |
| self._rope = (a.original_seq_len, a.compress_rope_theta) |
| else: |
| self._rope = (0, a.rope_theta) |
| self._rope_args = (a.rope_factor, a.beta_fast, a.beta_slow, a.max_seq_len) |
| self._fc = None |
|
|
| def freqs(self, seqlen, device): |
| if self._fc is None or self._fc.size(0) < seqlen or self._fc.device != device: |
| factor, bf, bs, max_len = self._rope_args |
| self._fc = precompute_freqs_cis(self.rd, max(seqlen, max_len), self._rope[0], self._rope[1], |
| factor, bf, bs, device) |
| return self._fc[:seqlen] |
|
|
| def forward(self, x, shared: Shared): |
| bsz, seqlen, _ = x.shape |
| fc = self.freqs(seqlen, x.device) |
| qr = self.q_norm(self.wq_a(x)) |
| q = rope(self.wq_b(qr).unflatten(-1, (self.n_heads, self.head_dim)), fc, self.rd) |
| kv = rope(self.kv_norm(self.wkv(x)), fc, self.rd) |
|
|
| comp_kv, cidx = None, None |
| if self.ratio: |
| latent = self.compressor(x) if self.is_kv_source else None |
| if self.is_index_source: |
| shared.topk_idxs = self.indexer(x, qr, latent, fc, shared, chunk=self.attn_chunk) |
| cidx = shared.topk_idxs |
| if latent is not None: |
| shared.compress_kv = rope(latent, fc[: seqlen - seqlen % self.ratio : self.ratio], self.rd) |
| comp_kv = shared.compress_kv |
|
|
| o = chunked_sink_attention(q, kv, self.attn_sink, self.window, self.softmax_scale, comp_kv, cidx, |
| chunk=self.attn_chunk) |
| o = rope(o, fc, self.rd, inverse=True) |
| o = o.reshape(bsz, seqlen, self.n_groups, -1) |
| wo_a = self.wo_a.weight.view(self.n_groups, self.o_lora_rank, -1) |
| o = torch.einsum("bsgd,grd->bsgr", o, wo_a.to(o.dtype)) |
| return self.wo_b(o.flatten(2)) |
|
|
|
|
| |
| |
|
|
|
|
| class Gate(nn.Module): |
| """sqrtsoftplus scores; a bias picks experts, weights come from the unbiased scores.""" |
|
|
| def __init__(self, a: Args): |
| super().__init__() |
| self.topk, self.score_func = a.n_activated_experts, a.score_func |
| self.norm_topk_prob, self.route_scale = a.norm_topk_prob, a.route_scale |
| self.weight = nn.Parameter(torch.empty(a.n_routed_experts, a.dim)) |
| self.register_buffer("bias", torch.zeros(a.n_routed_experts)) |
|
|
| def forward(self, x): |
| scores = F.linear(x.float(), self.weight.float()) |
| if self.score_func == "softmax": |
| scores = scores.softmax(dim=-1) |
| elif self.score_func == "sigmoid": |
| scores = scores.sigmoid() |
| else: |
| scores = F.softplus(scores).sqrt() |
| indices = (scores + self.bias.float()).topk(self.topk, dim=-1)[1] |
| weights = scores.gather(1, indices) |
| if self.norm_topk_prob and self.topk > 1: |
| weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-20) |
| return weights * self.route_scale, indices |
|
|
|
|
| class Expert(nn.Module): |
| def __init__(self, dim, inter_dim, swiglu_limit): |
| super().__init__() |
| self.w1 = nn.Linear(dim, inter_dim, bias=False) |
| self.w2 = nn.Linear(inter_dim, dim, bias=False) |
| self.w3 = nn.Linear(dim, inter_dim, bias=False) |
| self.swiglu_limit = swiglu_limit |
|
|
| def forward(self, x, weights=None): |
| dtype = x.dtype |
| gate, up = self.w1(x).float(), self.w3(x).float() |
| if self.swiglu_limit > 0: |
| up = up.clamp(-self.swiglu_limit, self.swiglu_limit) |
| gate = gate.clamp(max=self.swiglu_limit) |
| h = F.silu(gate) * up |
| if weights is not None: |
| h = weights * h |
| return self.w2(h.to(dtype)) |
|
|
|
|
| class MoE(nn.Module): |
| def __init__(self, a: Args): |
| super().__init__() |
| self.dim = a.dim |
| self.gate = Gate(a) |
| self.experts = nn.ModuleList([Expert(a.dim, a.moe_inter_dim, a.swiglu_limit) for _ in range(a.n_routed_experts)]) |
| assert a.n_shared_experts == 1 |
| self.shared_experts = Expert(a.dim, a.moe_inter_dim, a.swiglu_limit) |
|
|
| def forward(self, x): |
| shape = x.shape |
| x = x.reshape(-1, self.dim) |
| weights, indices = self.gate(x) |
| y = torch.zeros_like(x, dtype=torch.float32) |
| counts = torch.bincount(indices.flatten(), minlength=len(self.experts)).tolist() |
| for i, expert in enumerate(self.experts): |
| if counts[i] == 0: |
| continue |
| idx, top = torch.where(indices == i) |
| y = y.index_add(0, idx, expert(x[idx], weights[idx, top, None]).float()) |
| y = y + self.shared_experts(x).float() |
| return y.to(x.dtype).view(shape) |
|
|
|
|
| |
| |
|
|
|
|
| class Engram(nn.Module): |
| """Gated n-gram lookup written into each hc copy of the residual stream.""" |
|
|
| def __init__(self, a: Args, layer_id: int): |
| super().__init__() |
| self.layer_hash_index = tuple(a.engram_layer_ids).index(layer_id) |
| self.dim, self.hc_mult, self.eps = a.dim, a.hc_mult, a.norm_eps |
| self.embed = nn.Embedding(a.engram_num_embeddings[self.layer_hash_index], a.engram_head_dim) |
| n_hash_cols = (a.engram_max_ngram_size - 1) * a.engram_n_heads |
| self.wkv = nn.Linear(n_hash_cols * a.engram_head_dim, a.dim * (a.hc_mult + 1), bias=False) |
| self.q_weight = nn.Parameter(torch.ones(a.hc_mult, a.dim)) |
| self.k_weight = nn.Parameter(torch.ones(a.hc_mult, a.dim)) |
|
|
| def forward(self, x, hash_ids): |
| kv = self.wkv(self.embed(hash_ids).flatten(-2)) |
| key, value = kv.split([self.hc_mult * self.dim, self.dim], dim=-1) |
| key = key.float().unflatten(-1, (self.hc_mult, self.dim)) |
| h = x.float() |
| rstd = torch.rsqrt(h.square().mean(-1) + self.eps) * torch.rsqrt(key.square().mean(-1) + self.eps) |
| dot = (h * (self.q_weight.float() * self.k_weight.float()) * key).sum(-1) * rstd * self.dim**-0.5 |
| gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(1e-6).sqrt(), dot)) |
| return (h + gate.unsqueeze(-1) * value.float().unsqueeze(-2)).to(x.dtype) |
|
|
|
|
| class Block(nn.Module): |
| """Hyper-Connections block: the residual stream is carried as hc_mult copies.""" |
|
|
| def __init__(self, layer_id: int, a: Args): |
| super().__init__() |
| self.norm_eps = a.norm_eps |
| self.attn = Attention(layer_id, a) |
| self.ffn = MoE(a) |
| self.engram = Engram(a, layer_id) if layer_id in a.engram_layer_ids else None |
| self.attn_norm = RMSNorm(a.dim, a.norm_eps) |
| self.ffn_norm = RMSNorm(a.dim, a.norm_eps) |
| self.hc, self.hc_iters, self.hc_eps = a.hc_mult, a.hc_sinkhorn_iters, a.hc_eps |
| mix_hc = (2 + self.hc) * self.hc |
| self.hc_attn_fn = nn.Parameter(torch.empty(mix_hc, self.hc * a.dim)) |
| self.hc_ffn_fn = nn.Parameter(torch.empty(mix_hc, self.hc * a.dim)) |
| self.hc_attn_base = nn.Parameter(torch.empty(mix_hc)) |
| self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc)) |
| self.hc_attn_scale = nn.Parameter(torch.ones(3)) |
| self.hc_ffn_scale = nn.Parameter(torch.ones(3)) |
|
|
| def hc_mixes(self, x, fn, scale, base): |
| x = x.flatten(2).float() |
| mixes = F.linear(x, fn.float()) * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.norm_eps) |
| return hc_split_sinkhorn(mixes, scale.float(), base.float(), self.hc, self.hc_iters, self.hc_eps) |
|
|
| @staticmethod |
| def hc_pre(x, pre_mix): |
| return torch.sum(pre_mix.unsqueeze(-1) * x.float(), dim=2).to(x.dtype) |
|
|
| @staticmethod |
| def hc_post(x, residual, post, comb): |
| y = post.unsqueeze(-1) * x.unsqueeze(-2) + torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2) |
| return y.to(residual.dtype) |
|
|
| def forward(self, x, pre_mix, shared: Shared): |
| residual = x |
| attn_pre, attn_post, attn_comb = self.hc_mixes(x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base) |
| h = self.attn(self.attn_norm(self.hc_pre(x, pre_mix)), shared) |
| x = self.hc_post(h, residual, attn_post, attn_comb) |
| residual = x |
| ffn_pre, ffn_post, ffn_comb = self.hc_mixes(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base) |
| h = self.ffn(self.ffn_norm(self.hc_pre(x, attn_pre))) |
| x = self.hc_post(h, residual, ffn_post, ffn_comb) |
| return x, ffn_pre |
|
|
|
|
| class DeepseekV41PreTrainedModel(PreTrainedModel): |
| config_class = DeepseekV41Config |
| base_model_prefix = "" |
| supports_gradient_checkpointing = False |
|
|
| def _init_weights(self, module): |
| pass |
|
|
|
|
| class DeepseekV41ForCausalLM(DeepseekV41PreTrainedModel, GenerationMixin): |
| """Submodule names (embed, layers, norm, head) match the checkpoint's tensor names directly.""" |
|
|
| def __init__(self, config: DeepseekV41Config): |
| super().__init__(config) |
| a = Args.from_config(config) |
| self.args = a |
| self.embed = nn.Embedding(a.vocab_size, a.dim) |
| self.layers = nn.ModuleList([Block(i, a) for i in range(a.n_layers)]) |
| self.norm = RMSNorm(a.dim, a.norm_eps) |
| self.head = nn.Linear(a.dim, a.vocab_size, bias=False) |
| self.hc = a.hc_mult |
| self._engram_hash = None |
| self.post_init() |
|
|
| def get_input_embeddings(self): |
| return self.embed |
|
|
| def set_input_embeddings(self, value): |
| self.embed = value |
|
|
| def get_output_embeddings(self): |
| return self.head |
|
|
| def set_output_embeddings(self, value): |
| self.head = value |
|
|
| def _engram(self): |
| if self._engram_hash is None: |
| src = getattr(self.config, "_name_or_path", None) or self.config.name_or_path |
| self._engram_hash = NgramHash(self.args, AutoTokenizer.from_pretrained(src)) |
| return self._engram_hash |
|
|
| def forward(self, input_ids=None, attention_mask=None, labels=None, logits_to_keep=0, **kwargs): |
| """No KV cache; attention_mask is ignored (use one unpadded sequence per row, i.e. batch of equal |
| lengths or batch size 1). logits_to_keep=N returns logits for the last N positions only.""" |
| bsz, seqlen = input_ids.shape |
| hashes = self._engram()(input_ids) |
| h = self.embed(input_ids).unsqueeze(2).repeat(1, 1, self.hc, 1) |
| pre_mix = torch.zeros(bsz, seqlen, self.hc, device=h.device) |
| pre_mix[..., 0] = 1.0 |
| shared = Shared() |
| for layer in self.layers: |
| if layer.engram is not None: |
| h = layer.engram(h, hashes[:, :, layer.engram.layer_hash_index]) |
| h, pre_mix = layer(h, pre_mix, shared) |
| h = self.norm(Block.hc_pre(h, pre_mix)) |
| logits = self.head(h[:, -logits_to_keep:] if logits_to_keep else h) |
| loss = None |
| if labels is not None: |
| loss = F.cross_entropy(logits[:, :-1].float().flatten(0, 1), labels[:, 1:].flatten(), ignore_index=-100) |
| return CausalLMOutput(loss=loss, logits=logits) |
|
|
