| """Trainable PyTorch port of NVIDIA Dory1.0 (looped normalized-GPT). |
| |
| Mirrors vLLM's inference implementation (vllm-project/vllm#60149, |
| vllm/model_executor/models/dory.py) op for op, replacing its paged attention and |
| tensor-parallel layers with plain modules so the model can be fine-tuned: |
| |
| * full-sequence causal attention with per-layer sliding windows; each recurrent |
| loop attends over that loop's own keys/values, which is exactly vLLM's |
| "per_loop" KV-cache mode |
| * loop count selectable per forward pass (n_loops), for elastic-loop training |
| * per-layer gradient checkpointing and a chunked cross-entropy over the 131k vocab |
| * packed sequences (doc_ids + per-document position_ids); attention via dense-mask |
| SDPA (any device, used for checks) or FlexAttention block masks (CUDA training) |
| |
| Parameter names match the released checkpoint, so state dicts load and save as-is. |
| """ |
|
|
| import json |
| import math |
| from dataclasses import dataclass, fields |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
| from safetensors.torch import load_file |
| from torch import nn |
| from torch.utils.checkpoint import checkpoint |
|
|
| EPS = 1e-6 |
|
|
|
|
| @dataclass |
| class DoryConfig: |
| vocab_size: int = 131072 |
| hidden_size: int = 2560 |
| num_hidden_layers: int = 72 |
| intermediate_size: int = 7168 |
| num_attention_heads: int = 16 |
| num_key_value_heads: int = 4 |
| head_dim: int = 256 |
| initializer_range: float = 0.0177 |
| sliding_window: int = 1024 |
| rope_parameters: dict = None |
| layer_types: list = None |
| ngpt_alpha_init: float = 0.05 |
| attention_output_gate_logit_scale_init: float = None |
| n_input_layers: int = 16 |
| n_recurrent_layers: int = 40 |
| n_output_layers: int = 16 |
| n_recurrent_loops: int = 2 |
|
|
| @classmethod |
| def from_json(cls, path): |
| raw = json.loads(Path(path).read_text()) |
| known = {f.name for f in fields(cls)} |
| cfg = cls(**{k: v for k, v in raw.items() if k in known}) |
| assert cfg.n_input_layers + cfg.n_recurrent_layers + cfg.n_output_layers == cfg.num_hidden_layers |
| assert not raw.get("tie_word_embeddings") and raw.get("hidden_act", "silu") == "silu" |
| assert not raw.get("attention_bias") and not raw.get("mlp_bias") |
| return cfg |
|
|
|
|
| def justnorm(x, dim=-1): |
| """nGPT normalization: x * rsqrt(sum(x^2) + eps), fp32 accumulation.""" |
| xf = x.float() |
| return (xf * torch.rsqrt(xf.pow(2).sum(dim=dim, keepdim=True) + EPS)).to(x.dtype) |
|
|
|
|
| def slerp_alpha(a, b, alpha, alpha_scale): |
| ya, yb = justnorm(a).float(), justnorm(b).float() |
| lr = (alpha.float() * alpha_scale).abs() |
| return justnorm(ya + lr * (yb - ya)).to(a.dtype) |
|
|
|
|
| def rope_cos_sin(positions, head_dim, theta, device): |
| inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim)) |
| freqs = positions.to(device=device, dtype=torch.float32)[..., None] * inv_freq |
| return freqs.cos(), freqs.sin() |
|
|
|
|
| def apply_rope_neox(x, cos, sin): |
| """x: [B, T, H, D]; cos/sin: [B, T, D/2]. NeoX style (rotate halves).""" |
| x1, x2 = x.float().chunk(2, dim=-1) |
| c, s = cos[:, :, None, :], sin[:, :, None, :] |
| return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1).to(x.dtype) |
|
|
|
|
| def causal_mask(t, window, device, doc_ids=None): |
| """Bool mask, True = attend: [T, T], or [B, 1, T, T] with doc_ids (packed sequences). |
| Window w keeps keys with 0 <= i - j < w.""" |
| i = torch.arange(t, device=device)[:, None] |
| j = torch.arange(t, device=device)[None, :] |
| allowed = j <= i |
| if window is not None: |
| allowed &= (i - j) < window |
| if doc_ids is not None: |
| allowed = allowed[None] & (doc_ids[:, :, None] == doc_ids[:, None, :]) |
| allowed = allowed[:, None] |
| return allowed |
|
|
|
|
| _flex = None |
|
|
|
|
| def flex_block_mask(t, window, device, doc_ids): |
| from torch.nn.attention.flex_attention import create_block_mask |
|
|
| def mask_mod(b, h, q, kv): |
| ok = (kv <= q) & (doc_ids[b, q] == doc_ids[b, kv]) |
| if window is not None: |
| ok = ok & ((q - kv) < window) |
| return ok |
|
|
| return create_block_mask(mask_mod, doc_ids.shape[0], None, t, t, device=device) |
|
|
|
|
| def flex_attention(q, k, v, block_mask, scale): |
| global _flex |
| if _flex is None: |
| from torch.nn.attention.flex_attention import flex_attention as fa |
|
|
| _flex = torch.compile(fa, dynamic=False) |
| return _flex(q, k, v, block_mask=block_mask, scale=scale, enable_gqa=True) |
|
|
|
|
| class DoryAttention(nn.Module): |
| """GQA attention with nGPT qk-norm (after RoPE), sigmoid output gate and justnorm.""" |
|
|
| def __init__(self, cfg: DoryConfig, layer_type: str): |
| super().__init__() |
| h, kv, d = cfg.num_attention_heads, cfg.num_key_value_heads, cfg.head_dim |
| self.num_heads, self.num_kv_heads, self.head_dim = h, kv, d |
| self.q_proj = nn.Linear(cfg.hidden_size, h * d, bias=False) |
| self.k_proj = nn.Linear(cfg.hidden_size, kv * d, bias=False) |
| self.v_proj = nn.Linear(cfg.hidden_size, kv * d, bias=False) |
| self.g_proj = nn.Linear(cfg.hidden_size, h * d, bias=False) |
| self.o_proj = nn.Linear(h * d, cfg.hidden_size, bias=False) |
| self.sqk = nn.Parameter(torch.empty(h * d)) |
| self.s_gate = nn.Parameter(torch.empty(h * d)) |
| std = cfg.initializer_range |
| self.sqk_factor = d**0.5 / std |
| g0 = cfg.attention_output_gate_logit_scale_init or cfg.hidden_size**0.5 |
| self.s_gate_factor = g0**0.5 / std |
| self.layer_type = layer_type |
|
|
| def forward(self, x, cos, sin, mask): |
| b, t, _ = x.shape |
| h, kv, d = self.num_heads, self.num_kv_heads, self.head_dim |
| q = self.q_proj(x).view(b, t, h, d) |
| k = self.k_proj(x).view(b, t, kv, d) |
| v = self.v_proj(x).view(b, t, kv, d) |
| gate = self.g_proj(x) |
|
|
| q, k = apply_rope_neox(q, cos, sin), apply_rope_neox(k, cos, sin) |
| sqk = (self.sqk.float() * self.sqk_factor).view(h, d) |
| q = (sqk.square() * justnorm(q).float()).to(q.dtype) |
| k = justnorm(k) |
|
|
| q, k, v = (z.transpose(1, 2) for z in (q, k, v)) |
| if isinstance(mask, torch.Tensor): |
| if kv != h: |
| k = k.repeat_interleave(h // kv, dim=1) |
| v = v.repeat_interleave(h // kv, dim=1) |
| out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=1.0 / math.sqrt(d)) |
| else: |
| out = flex_attention(q, k, v, mask, 1.0 / math.sqrt(d)) |
| out = out.transpose(1, 2).reshape(b, t, h * d) |
|
|
| sg = (self.s_gate.float() * self.s_gate_factor).square() |
| out = (out.float() * torch.sigmoid(gate.float() * sg)).to(out.dtype) |
| return self.o_proj(justnorm(out)) |
|
|
|
|
| class DoryMLP(nn.Module): |
| """SwiGLU with nGPT suv scaling and justnorm before down_proj.""" |
|
|
| def __init__(self, cfg: DoryConfig): |
| super().__init__() |
| self.gate_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False) |
| self.up_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False) |
| self.down_proj = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False) |
| self.suv_gate = nn.Parameter(torch.empty(cfg.intermediate_size)) |
| self.suv_up = nn.Parameter(torch.empty(cfg.intermediate_size)) |
| self.suv_factor = cfg.hidden_size**0.5 |
|
|
| def forward(self, x): |
| g = (self.gate_proj(x).float() * (self.suv_gate.float() * self.suv_factor)).to(x.dtype) |
| u = (self.up_proj(x).float() * (self.suv_up.float() * self.suv_factor)).to(x.dtype) |
| m = (F.silu(g.float()) * u.float()).to(x.dtype) |
| return self.down_proj(justnorm(m)) |
|
|
|
|
| class DoryLayer(nn.Module): |
| """One physical layer: a single sublayer (attention OR mlp) plus slerp residual.""" |
|
|
| def __init__(self, cfg: DoryConfig, idx: int): |
| super().__init__() |
| self.layer_type = cfg.layer_types[idx] |
| self.is_attention = self.layer_type != "mlp" |
| self.alpha_scale = cfg.ngpt_alpha_init / cfg.initializer_range |
| if self.is_attention: |
| self.mixer = DoryAttention(cfg, self.layer_type) |
| self.attn_alpha = nn.Parameter(torch.empty(cfg.hidden_size)) |
| else: |
| self.mixer = DoryMLP(cfg) |
| self.mlp_alpha = nn.Parameter(torch.empty(cfg.hidden_size)) |
|
|
| def forward(self, h, ctx): |
| if self.is_attention: |
| cos, sin = ctx["rope"][self.layer_type] |
| out = self.mixer(h, cos, sin, ctx["mask"][self.layer_type]) |
| return slerp_alpha(h, out, self.attn_alpha, self.alpha_scale) |
| return slerp_alpha(h, self.mixer(h), self.mlp_alpha, self.alpha_scale) |
|
|
|
|
| class DoryBackbone(nn.Module): |
| def __init__(self, cfg: DoryConfig): |
| super().__init__() |
| self.cfg = cfg |
| self.embeddings = nn.Embedding(cfg.vocab_size, cfg.hidden_size) |
| self.layers = nn.ModuleList(DoryLayer(cfg, i) for i in range(cfg.num_hidden_layers)) |
| self.gradient_checkpointing = False |
|
|
| def _run(self, layer, h, ctx): |
| if self.gradient_checkpointing and self.training and torch.is_grad_enabled(): |
| return checkpoint(layer, h, ctx, use_reentrant=False) |
| return layer(h, ctx) |
|
|
| def forward(self, input_ids, position_ids=None, n_loops=None, doc_ids=None, attn_impl="sdpa"): |
| cfg = self.cfg |
| b, t = input_ids.shape |
| dev = input_ids.device |
| if position_ids is None: |
| position_ids = torch.arange(t, device=dev).expand(b, t) |
| if attn_impl == "flex" and doc_ids is None: |
| doc_ids = torch.zeros(b, t, dtype=torch.int32, device=dev) |
| ctx = {"rope": {}, "mask": {}} |
| for lt in set(cfg.layer_types) - {"mlp"}: |
| ctx["rope"][lt] = rope_cos_sin(position_ids, cfg.head_dim, cfg.rope_parameters[lt]["rope_theta"], dev) |
| window = cfg.sliding_window if lt == "sliding_attention" else None |
| if attn_impl == "flex": |
| ctx["mask"][lt] = flex_block_mask(t, window, dev, doc_ids) |
| else: |
| ctx["mask"][lt] = causal_mask(t, window, dev, doc_ids) |
|
|
| n_in, n_rec = cfg.n_input_layers, cfg.n_recurrent_layers |
| h = self.embeddings(input_ids) |
| for layer in self.layers[:n_in]: |
| h = self._run(layer, h, ctx) |
| input_h = h |
| state = input_h |
| for _ in range(n_loops or cfg.n_recurrent_loops): |
| h = justnorm(state + input_h) |
| for layer in self.layers[n_in : n_in + n_rec]: |
| h = self._run(layer, h, ctx) |
| state = justnorm(h) |
| h = state |
| for layer in self.layers[n_in + n_rec :]: |
| h = self._run(layer, h, ctx) |
| return h |
|
|
|
|
| class DoryForCausalLM(nn.Module): |
| def __init__(self, cfg: DoryConfig): |
| super().__init__() |
| self.config = cfg |
| self.backbone = DoryBackbone(cfg) |
| self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False) |
| self.lm_head.logit_scale = nn.Parameter(torch.empty(cfg.vocab_size)) |
|
|
| @classmethod |
| def from_pretrained(cls, path, dtype=torch.bfloat16, device="cpu"): |
| path = Path(path) |
| cfg = DoryConfig.from_json(path / "config.json") |
| with torch.device("meta"): |
| model = cls(cfg) |
| state = {} |
| for shard in sorted(path.glob("*.safetensors")): |
| state.update(load_file(shard)) |
| state = {k: v.to(dtype) for k, v in state.items()} |
| missing, unexpected = model.load_state_dict(state, strict=True, assign=True) |
| assert not missing and not unexpected, (missing, unexpected) |
| return model.to(device) |
|
|
| def gradient_checkpointing_enable(self): |
| self.backbone.gradient_checkpointing = True |
|
|
| @torch.no_grad() |
| def normalize_matrices_(self): |
| """Project weights back onto the unit hypersphere after an optimizer step (nGPT). |
| |
| The released checkpoint has every matrix unit-norm along the hidden dimension: |
| rows of embeddings/lm_head and of maps out of the residual stream (q/k/v/g, |
| gate/up), columns of maps into it (o_proj, down_proj). |
| """ |
| hidden = self.config.hidden_size |
| for w in (self.backbone.embeddings.weight, self.lm_head.weight): |
| w.copy_(justnorm(w, dim=1)) |
| for mod in self.backbone.modules(): |
| if isinstance(mod, nn.Linear): |
| dim = 1 if mod.in_features == hidden else 0 |
| mod.weight.copy_(justnorm(mod.weight, dim=dim)) |
|
|
| def logits(self, hidden): |
| scale = self.lm_head.logit_scale.float() / self.config.initializer_range |
| return self.lm_head(hidden).float() * scale |
|
|
| def forward(self, input_ids, labels=None, position_ids=None, n_loops=None, doc_ids=None, |
| attn_impl="sdpa", loss_chunk=2048): |
| """Without labels: returns float32 logits [B, T, V]. |
| With labels (same shape as input_ids, -100 = ignore; shifted internally, so a |
| packed batch must set -100 on each document's first token): |
| returns (mean_loss, n_target_tokens) via a chunked, checkpointed CE.""" |
| hidden = self.backbone(input_ids, position_ids=position_ids, n_loops=n_loops, |
| doc_ids=doc_ids, attn_impl=attn_impl) |
| if labels is None: |
| return self.logits(hidden) |
| h = hidden[:, :-1].reshape(-1, hidden.shape[-1]) |
| y = labels[:, 1:].reshape(-1) |
| keep = y != -100 |
| h, y = h[keep], y[keep] |
| total = h.new_zeros((), dtype=torch.float32) |
|
|
| def chunk_loss(hc, yc): |
| return F.cross_entropy(self.logits(hc), yc, reduction="sum") |
|
|
| for s in range(0, h.shape[0], loss_chunk): |
| hc, yc = h[s : s + loss_chunk], y[s : s + loss_chunk] |
| if self.training and torch.is_grad_enabled(): |
| total = total + checkpoint(chunk_loss, hc, yc, use_reentrant=False) |
| else: |
| total = total + chunk_loss(hc, yc) |
| n = int(keep.sum()) |
| return total / max(n, 1), n |
|
|
| @torch.no_grad() |
| def generate(self, input_ids, max_new_tokens=128, n_loops=None, temperature=0.0, eos_ids=()): |
| """Simple cache-free sampling, for correctness checks (not speed).""" |
| ids = input_ids |
| for _ in range(max_new_tokens): |
| logits = self.forward(ids, n_loops=n_loops)[:, -1] |
| if temperature > 0: |
| nxt = torch.multinomial(torch.softmax(logits / temperature, -1), 1) |
| else: |
| nxt = logits.argmax(-1, keepdim=True) |
| ids = torch.cat([ids, nxt], dim=1) |
| if int(nxt[0, 0]) in eos_ids: |
| break |
| return ids[:, input_ids.shape[1] :] |
|
|
