Hugging Face's logo

rdaum
/
decisive-dory

Text Generation
Safetensors
GGUF
English
dory
decision-making
lora
eider
bf16
nvfp4
llama-cpp
conversational
decision-model
Model card Files
xet
 Community
decisive-dory / adapter /dory_lora.py
rdaum's picture
rdaum
Publish state-32k BF16, NVFP4, and Q8_0 GGUF with shared model card
5e71a7a verified
Raw History Blame Contribute Delete
11.5 kB
"""BF16 Dory with frozen checkpoint weights and trainable linear LoRA adapters.
The nGPT equations follow the independent dory-checkpoint-reference.py oracle
and vLLM revision 55076950b6c155f384d6766d63563c3bc20a0e3f. This is a training
experiment, not a replacement for Eider's CUDA serving implementation.
"""
from __future__ import annotations
import json
import math
from pathlib import Path
import torch
from safetensors import safe_open
from safetensors.torch import load_file, save_file
from torch import nn
from torch.nn import functional as F
from torch.utils.checkpoint import checkpoint
def norm(x):
value = x.float()
return (value * torch.rsqrt(value.square().sum(-1, keepdim=True) + 1e-6)).to(
x.dtype
)
class Linear(nn.Module):
def __init__(self, weight, rank, alpha):
super().__init__()
self.register_buffer("weight", weight)
self.scale = alpha / rank if rank else 0
if rank:
self.lora_a = nn.Parameter(
torch.empty(rank, weight.shape[1], device=weight.device)
)
self.lora_b = nn.Parameter(
torch.zeros(weight.shape[0], rank, device=weight.device)
)
nn.init.kaiming_uniform_(self.lora_a, a=math.sqrt(5))
else:
self.register_parameter("lora_a", None)
self.register_parameter("lora_b", None)
def forward(self, x):
base = F.linear(x, self.weight)
if self.lora_a is None:
return base
delta = F.linear(F.linear(x, self.lora_a.to(x.dtype)), self.lora_b.to(x.dtype))
return (base.float() + self.scale * delta.float()).to(x.dtype)
class Layer(nn.Module):
def __init__(self, config, index, read, rank, alpha):
super().__init__()
self.config = config
self.kind = config["layer_types"][index]
prefix = f"backbone.layers.{index}"
mixer = prefix + ".mixer"
std = config["initializer_range"]
if self.kind == "mlp":
names = ("gate_proj", "up_proj", "down_proj")
scale = math.sqrt(config["hidden_size"])
self.register_buffer("suv_gate", read(mixer + ".suv_gate").float() * scale)
self.register_buffer("suv_up", read(mixer + ".suv_up").float() * scale)
rate = read(prefix + ".mlp_alpha")
else:
names = ("q_proj", "k_proj", "v_proj", "g_proj", "o_proj")
dim = config["head_dim"]
self.register_buffer(
"sqk", (read(mixer + ".sqk").float() * (math.sqrt(dim) / std)).square()
)
g0 = config.get("attention_output_gate_logit_scale_init") or math.sqrt(
config["hidden_size"]
)
self.register_buffer(
"s_gate",
(read(mixer + ".s_gate").float() * (math.sqrt(g0) / std)).square(),
)
rate = read(prefix + ".attn_alpha")
self.register_buffer(
"rate", (rate.float() * config["ngpt_alpha_init"] / std).abs()
)
self.projections = nn.ModuleDict(
{
name: Linear(read(mixer + "." + name + ".weight"), rank, alpha)
for name in names
}
)
def forward(self, h, rotary, sliding_mask, reference=False):
p, c, dtype = self.projections, self.config, h.dtype
if self.kind == "mlp":
gate = (p["gate_proj"](h).float() * self.suv_gate).to(dtype)
up = (p["up_proj"](h).float() * self.suv_up).to(dtype)
out = p["down_proj"](norm((F.silu(gate.float()) * up.float()).to(dtype)))
else:
nq, nk, dim = (
c["num_attention_heads"],
c["num_key_value_heads"],
c["head_dim"],
)
cos, sin = rotary[self.kind]
def rope(x):
a, b = x.float().chunk(2, -1)
return torch.cat((a * cos - b * sin, b * cos + a * sin), -1).to(dtype)
q = norm(rope(p["q_proj"](h).reshape(*h.shape[:-1], nq, dim)))
k = norm(rope(p["k_proj"](h).reshape(*h.shape[:-1], nk, dim)))
q = (q.float() * self.sqk.reshape(nq, dim)).to(dtype)
v = p["v_proj"](h).reshape(*h.shape[:-1], nk, dim)
# Repeat GQA heads explicitly so both SDPA backends and the oracle
# use the same layout. Batch rows retain independent attention.
q = q.transpose(1, 2)
k = k.repeat_interleave(nq // nk, 2).transpose(1, 2)
v = v.repeat_interleave(nq // nk, 2).transpose(1, 2)
mask = sliding_mask if self.kind == "sliding_attention" else None
if reference:
if mask is None:
mask = torch.ones(
h.shape[1], h.shape[1], device=h.device, dtype=torch.bool
).tril()
scores = (q.float() @ k.float().transpose(-1, -2)) / math.sqrt(dim)
out = (
scores.masked_fill(~mask, -torch.inf).softmax(-1) @ v.float()
).to(dtype)
else:
out = F.scaled_dot_product_attention(
q, k, v, attn_mask=mask, is_causal=mask is None
)
out = out.transpose(1, 2).reshape(*h.shape[:-1], nq * dim)
gate = p["g_proj"](h)
out = norm(
(out.float() * torch.sigmoid(gate.float() * self.s_gate)).to(dtype)
)
out = p["o_proj"](out)
a, b = norm(h).float(), norm(out).float()
return norm(a + self.rate * (b - a)).to(dtype)
class Dory(nn.Module):
def __init__(self, config, read, *, rank=8, alpha=16, loops=2, checkpointing=False):
super().__init__()
if config.get("model_type") != "dory" or config.get("quantization_config"):
raise ValueError("requires an unquantized Dory checkpoint")
if loops < 1 or rank < 0 or alpha <= 0:
raise ValueError("invalid loop count or adapter dimensions")
self.config, self.loops, self.checkpointing = config, loops, checkpointing
self.register_buffer("embedding", read("backbone.embeddings.weight"))
self.register_buffer("head", read("lm_head.weight"))
self.register_buffer(
"head_scale",
read("lm_head.logit_scale").float() / config["initializer_range"],
)
self.layers = nn.ModuleList(
Layer(config, i, read, rank, alpha)
for i in range(config["num_hidden_layers"])
)
@classmethod
def load(cls, path: Path, *, device="cuda", **kwargs):
config = json.loads((path / "config.json").read_text())
with safe_open(
path / "model.safetensors", framework="pt", device="cpu"
) as weights:
consumed = set()
def read(name):
value = weights.get_tensor(name)
if value.dtype != torch.bfloat16:
raise ValueError(f"expected BF16 weight: {name}")
consumed.add(name)
return value.to(device)
model = cls(config, read, **kwargs)
if consumed != set(weights.keys()):
raise ValueError(
f"unconsumed checkpoint tensors: {set(weights.keys()) - consumed}"
)
return model
def forward(self, ids, labels=None, *, lengths=None, reference=False):
c = self.config
single = ids.ndim == 1
if single:
ids = ids.unsqueeze(0)
positions = torch.arange(ids.shape[1], device=ids.device)
dim = c["head_dim"]
rotary = {}
for kind, params in c["rope_parameters"].items():
angles = positions.float()[:, None] * params["rope_theta"] ** (
-torch.arange(dim // 2, device=ids.device).float() / (dim // 2)
)
rotary[kind] = (angles.cos()[:, None, :], angles.sin()[:, None, :])
# For short training rows the sliding mask equals the causal mask.
sliding_mask = None
if ids.shape[1] > c["sliding_window"]:
distance = positions[:, None] - positions[None, :]
sliding_mask = (distance >= 0) & (distance < c["sliding_window"])
def apply(h, start, end):
for layer in self.layers[start:end]:
if self.training and self.checkpointing and not reference:
h = checkpoint(layer, h, rotary, sliding_mask, use_reentrant=False)
else:
# Keep the serial evaluator and oracle eager. Their batch-one,
# inference-only graphs otherwise compete with training graphs.
forward = (
layer if self.training and not reference else layer.forward
)
h = forward(h, rotary, sliding_mask, reference)
return h
ni, nr = c["n_input_layers"], c["n_recurrent_layers"]
h = apply(F.embedding(ids, self.embedding), 0, ni)
input_h = h
for _ in range(self.loops):
h = norm(h + input_h)
h = norm(apply(h, ni, ni + nr))
h = apply(h, ni + nr, len(self.layers))
# Right-padding is always after each real query, so causal attention
# excludes it without a padding mask. Read only each row's final token.
last = (
h[:, -1]
if lengths is None
else h[torch.arange(ids.shape[0], device=ids.device), lengths - 1]
)
weight = self.head if labels is None else self.head[labels]
scale = self.head_scale if labels is None else self.head_scale[labels]
logits = (
torch.bmm(weight, last.unsqueeze(-1)).squeeze(-1)
if labels is not None and labels.ndim == 2
else F.linear(last, weight)
).float() * scale
return logits.squeeze(0) if single else logits
def compile_layers(self):
# Compile individual layers so recurrent traversals reuse their graphs.
# Keep the BF16 rounding points used by the eager nGPT equations.
# All layer kinds share Layer.forward's code cache. Attention masks,
# input gradients and shape guards produce more than eight variants when
# training batches cross the sliding-window boundary. Keep a bounded
# cache large enough for those variants; fullgraph raises on overflow.
torch._dynamo.config.recompile_limit = max(
torch._dynamo.config.recompile_limit, 32
)
for layer in self.layers:
layer.compile(
fullgraph=True,
dynamic=True,
options={"emulate_precision_casts": True},
)
def save_adapter(self, path):
save_file(
{
name: p.detach().cpu().contiguous()
for name, p in self.named_parameters()
},
str(path),
)
def load_adapter(self, path):
values = load_file(str(path))
params = dict(self.named_parameters())
if values.keys() != params.keys():
raise ValueError("adapter does not match the model's projection modules")
with torch.no_grad():
for name, p in params.items():
if p.shape != values[name].shape:
raise ValueError(f"adapter shape mismatch: {name}")
p.copy_(values[name])