mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 07:17:08 +08:00
8a5209ee65
- examples/neuron: tiny random V4.1 checkpoint in the real HF layout, pure-torch replacements for the official reference kernels, golden generation, a teacher-forced DeepInfer runner (LLM offline API) and a comparison tool. CPU mode: 25/27 greedy tokens match, worst logits cos 0.9985 / rel 0.056 against a reference self-noise floor of 0.032-0.060. - model_loader: check parameters against the platform device, not is_cuda. - deep_gemm paged MQA logits: do not dereference block-table slots past the context. - bind_cpu: skip in Neuron CPU mode (1 core per NeuronCore share would serialize host compute). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
466 lines
21 KiB
Python
466 lines
21 KiB
Python
"""Generate a tiny, random, deterministic DeepSeek-V4.1-Flash HF checkpoint.
|
|
|
|
Same tensor names, dtypes and layout conventions as the released checkpoint
|
|
(deepseek-ai/DeepSeek-V4.1-Flash-1017); only the sizes shrink. The result is meant to be loadable by
|
|
both DeepInfer's HF loader (deepinfer/hugging_face/*) and the official reference
|
|
(inference/convert.py + inference/model.py).
|
|
|
|
Conventions (taken from the official inference/ code):
|
|
* FP8 weights: F8_E4M3 [out, in] + F8_E8M0 scale [ceil(out/32), ceil(in/32)] (32x32 blocks).
|
|
scale = 2**ceil(log2(amax/448)), q = e4m3(clamp(w/scale, +-448)); dequant = q * scale.
|
|
* Routed experts: FP4 e2m1, two codes per byte stored as I8 [out, in/2]; element 2j is the LOW nibble,
|
|
element 2j+1 the HIGH nibble (convert.py cast_e2m1fn_to_e4m3fn: low = x & 0xF -> first, high -> second);
|
|
code table = FP4_TABLE in convert.py (bit 3 = sign, 0..7 = 0,.5,1,1.5,2,3,4,6).
|
|
F8_E8M0 scale [out, in/32] (1x32 blocks along K), scale = 2**ceil(log2(amax/6)), RNE to the e2m1 grid.
|
|
* Engram table: F8_E4M3 [rows, head_dim] + F8_E8M0 [rows, head_dim/32] (1x32 blocks), rows =
|
|
sum of the per-(ngram, head) prime bucket sizes (EngramLayout in DeepInfer and in inference/engram.py).
|
|
* hc_* and attn_sink and gate biases are F32; norms, compressor, indexer wk/weights_proj, gate weight,
|
|
embed/head, vision, markov/confidence heads are BF16 (as in the real index).
|
|
|
|
Random weights: W ~ N(0, 1/in_features) in fp32 -> bf16 (-> quantized), norms = 1, biases/sinks small; the
|
|
projections that write into the residual stream get an extra gain (attn wo_b: --attn-out-scale 0.5; shared/routed
|
|
expert w2 and engram wkv: --ffn-out-scale 1/16) so that the reference is numerically stable enough to be a golden.
|
|
Every tensor uses its own generator seeded from (seed, crc32(name)), so the content of one tensor does not
|
|
depend on which other tensors exist.
|
|
|
|
Usage:
|
|
python make_tiny_checkpoint.py --out ~/di/tiny_ckpt --tokenizer-dir <dir with tokenizer.json> \
|
|
--real-config <real HF config.json>
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import copy
|
|
import json
|
|
import math
|
|
import os
|
|
import shutil
|
|
import zlib
|
|
from collections import OrderedDict
|
|
|
|
import torch
|
|
from safetensors.torch import save_file
|
|
|
|
# ----------------------------------------------------------------------------------------------
|
|
# Tiny config. Keys not listed keep the value from the real HF config.
|
|
# ----------------------------------------------------------------------------------------------
|
|
TINY_TEXT_OVERRIDES = dict(
|
|
hidden_size=1024,
|
|
moe_intermediate_size=256,
|
|
num_hidden_layers=10,
|
|
num_attention_heads=8,
|
|
head_dim=512,
|
|
qk_rope_head_dim=64,
|
|
q_lora_rank=256,
|
|
o_lora_rank=256,
|
|
o_groups=2,
|
|
n_routed_experts=16,
|
|
num_experts_per_tok=4,
|
|
# 10 backbone layers + 3 nextn layers. Layout mirrors the real one: 2 SWA-only layers, a ratio-2
|
|
# segment, a ratio-1 segment starting at the candidate source, 3 SWA-only MTP layers.
|
|
compress_ratios=[0, 0, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0],
|
|
kv_source_layer_ids=[2, 4, 6],
|
|
index_source_layer_ids=[2, 4, 6, 8],
|
|
candidate_source_layer_id=6,
|
|
# DeepInfer: candidate_tokens * sizeof(bf16 psum) % 1024 == 0 (index_attn.py:109) -> topk_blocks*8 % 512 == 0,
|
|
# so 64 is the smallest legal value with block size 8.
|
|
candidate_topk_blocks=64,
|
|
index_n_heads=4,
|
|
index_head_dim=128,
|
|
index_topk=64,
|
|
engram_layer_ids=[1, 4],
|
|
engram_vocab_size=4096,
|
|
engram_n_heads=2,
|
|
engram_head_dim=256,
|
|
num_nextn_predict_layers=3,
|
|
dspark_target_layer_ids=[7, 8, 9],
|
|
dspark_n_routed_experts=8,
|
|
dspark_num_experts_per_tok=2,
|
|
)
|
|
# DeepInfer's VisionTransformer hard-codes DEPTH=32 / dim 1024 / 16 heads / mlp 2816
|
|
# (deepinfer/model_executor/models/vision_encoder/vit.py), always builds it, and the loader requires every
|
|
# model parameter to be covered (model_loader.load_hf_checkpoint). So the vision tower keeps its real size.
|
|
TINY_VISION_OVERRIDES = dict()
|
|
|
|
FP8_MAX = 448.0
|
|
FP4_MAX = 6.0
|
|
BLOCK = 32
|
|
FP4_TABLE = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0])
|
|
|
|
|
|
# ----------------------------------------------------------------------------------------------
|
|
# Quantization helpers (match inference/kernel.py fast_round_scale: 2**ceil(log2(amax * (1/max))))
|
|
# ----------------------------------------------------------------------------------------------
|
|
def pow2_ceil_scale(amax: torch.Tensor, max_inv: float) -> torch.Tensor:
|
|
"""2**ceil(log2(amax * max_inv)) computed from the fp32 bit pattern, like kernel.py fast_log2_ceil."""
|
|
v = (amax.float() * torch.tensor(max_inv, dtype=torch.float32)).contiguous()
|
|
bits = v.view(torch.int32)
|
|
exp = (bits >> 23) & 0xFF
|
|
man = bits & ((1 << 23) - 1)
|
|
e = exp - 127 + (man != 0).to(torch.int32)
|
|
return torch.ldexp(torch.ones_like(v), e)
|
|
|
|
|
|
def to_e8m0(scale: torch.Tensor) -> torch.Tensor:
|
|
"""Power-of-two fp32 -> float8_e8m0fnu (exact)."""
|
|
bits = (scale.float().contiguous().view(torch.int32) >> 23) & 0xFF
|
|
return bits.to(torch.uint8).view(torch.float8_e8m0fnu)
|
|
|
|
|
|
def quant_fp8_block(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""[out, in] -> e4m3 weight + e8m0 scale [ceil(out/32), ceil(in/32)]."""
|
|
out_f, in_f = w.shape
|
|
bo, bi = math.ceil(out_f / BLOCK), math.ceil(in_f / BLOCK)
|
|
wp = torch.zeros(bo * BLOCK, bi * BLOCK, dtype=torch.float32)
|
|
wp[:out_f, :in_f] = w.float()
|
|
blocks = wp.view(bo, BLOCK, bi, BLOCK)
|
|
amax = blocks.abs().amax(dim=(1, 3)).clamp_min(1e-4)
|
|
s = pow2_ceil_scale(amax, 1.0 / FP8_MAX)
|
|
q = (blocks / s[:, None, :, None]).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
|
|
q = q.view(bo * BLOCK, bi * BLOCK)[:out_f, :in_f].contiguous()
|
|
return q, to_e8m0(s)
|
|
|
|
|
|
def round_e2m1_codes(v: torch.Tensor) -> torch.Tensor:
|
|
"""fp32 (already scaled, |v|<=6) -> 4-bit e2m1 codes, round-to-nearest-even."""
|
|
a = v.abs()
|
|
code = torch.zeros_like(a, dtype=torch.int32)
|
|
# midpoints whose tie goes DOWN to an even code: 0.25(0|1) 1.25(2|3) 2.5(4|5) 5.0(6|7)
|
|
for b in (0.25, 1.25, 2.5, 5.0):
|
|
code += (a > b).to(torch.int32)
|
|
# midpoints whose tie goes UP to an even code: 0.75(1|2) 1.75(3|4) 3.5(5|6)
|
|
for b in (0.75, 1.75, 3.5):
|
|
code += (a >= b).to(torch.int32)
|
|
sign = (v < 0) | ((v == 0) & torch.signbit(v))
|
|
return code | (sign.to(torch.int32) << 3)
|
|
|
|
|
|
def quant_fp4_row(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""[out, in] -> packed int8 [out, in/2] (low nibble = even k) + e8m0 scale [out, in/32]."""
|
|
out_f, in_f = w.shape
|
|
assert in_f % BLOCK == 0
|
|
blocks = w.float().view(out_f, in_f // BLOCK, BLOCK)
|
|
amax = blocks.abs().amax(-1).clamp_min(FP4_MAX * 2.0**-126)
|
|
s = pow2_ceil_scale(amax, 1.0 / FP4_MAX)
|
|
codes = round_e2m1_codes((blocks / s[..., None]).clamp(-FP4_MAX, FP4_MAX)).view(out_f, in_f)
|
|
packed = (codes[:, 0::2] | (codes[:, 1::2] << 4)).to(torch.uint8).view(torch.int8).contiguous()
|
|
return packed, to_e8m0(s)
|
|
|
|
|
|
def dequant_fp4_row(packed: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
|
|
x = packed.view(torch.uint8)
|
|
lo, hi = (x & 0xF).long(), (x >> 4).long()
|
|
vals = torch.stack([FP4_TABLE[lo], FP4_TABLE[hi]], dim=-1).flatten(-2)
|
|
return vals * scale.float().repeat_interleave(BLOCK, dim=-1)
|
|
|
|
|
|
def quant_fp8_row(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Engram table: [rows, d] -> e4m3 + e8m0 [rows, d/32] (1x32 blocks)."""
|
|
rows, d = w.shape
|
|
blocks = w.float().view(rows, d // BLOCK, BLOCK)
|
|
amax = blocks.abs().amax(-1).clamp_min(1e-4)
|
|
s = pow2_ceil_scale(amax, 1.0 / FP8_MAX)
|
|
q = (blocks / s[..., None]).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn).view(rows, d)
|
|
return q.contiguous(), to_e8m0(s)
|
|
|
|
|
|
# ----------------------------------------------------------------------------------------------
|
|
# Engram layout (identical in DeepInfer engram_metadata.EngramLayout.from_cfg and inference/engram.py)
|
|
# ----------------------------------------------------------------------------------------------
|
|
def engram_primes(layer_ids, max_ngram_size, n_heads, vocab_size):
|
|
from sympy import isprime, nextprime
|
|
|
|
# reference (inference/engram.py find_next_prime)
|
|
seen, ref = set(), []
|
|
for _ in layer_ids:
|
|
per_ngram = []
|
|
for _ in range(max_ngram_size - 1):
|
|
sizes, current = [], vocab_size - 1
|
|
for _ in range(n_heads):
|
|
c = current + 1
|
|
while not isprime(c) or c in seen:
|
|
c += 1
|
|
current = c
|
|
seen.add(c)
|
|
sizes.append(c)
|
|
per_ngram.append(sizes)
|
|
ref.append(per_ngram)
|
|
# DeepInfer (engram_metadata.EngramLayout.from_cfg)
|
|
seen, di = set(), []
|
|
for _ in layer_ids:
|
|
per_ngram = []
|
|
for _ in range(max_ngram_size - 1):
|
|
start, sizes = vocab_size - 1, []
|
|
for _ in range(n_heads):
|
|
p = int(nextprime(start))
|
|
while p in seen:
|
|
p = int(nextprime(p))
|
|
seen.add(p)
|
|
sizes.append(p)
|
|
start = p
|
|
per_ngram.append(sizes)
|
|
di.append(per_ngram)
|
|
assert ref == di, (ref, di)
|
|
return ref
|
|
|
|
|
|
# ----------------------------------------------------------------------------------------------
|
|
# Tensor factory
|
|
# ----------------------------------------------------------------------------------------------
|
|
class Factory:
|
|
def __init__(self, seed: int, norm_noise: float):
|
|
self.seed = seed
|
|
self.norm_noise = norm_noise
|
|
self.tensors: "OrderedDict[str, torch.Tensor]" = OrderedDict()
|
|
|
|
def gen(self, name: str) -> torch.Generator:
|
|
g = torch.Generator()
|
|
g.manual_seed((self.seed * 1_000_003 + zlib.crc32(name.encode())) % (2**63))
|
|
return g
|
|
|
|
def randn(self, name, shape, std):
|
|
return torch.randn(shape, generator=self.gen(name), dtype=torch.float32) * std
|
|
|
|
def add(self, name, t):
|
|
assert name not in self.tensors, name
|
|
self.tensors[name] = t.contiguous()
|
|
|
|
def bf16(self, name, out_f, in_f, std=None):
|
|
std = (1.0 / in_f) ** 0.5 if std is None else std
|
|
self.add(name, self.randn(name, (out_f, in_f), std).bfloat16())
|
|
|
|
def norm(self, name, dim):
|
|
w = torch.ones(dim)
|
|
if self.norm_noise:
|
|
w = w + self.randn(name, (dim,), self.norm_noise)
|
|
self.add(name, w.bfloat16())
|
|
|
|
def bias(self, name, dim, std=0.02, dtype=torch.bfloat16):
|
|
self.add(name, self.randn(name, (dim,), std).to(dtype))
|
|
|
|
def fp8(self, prefix, out_f, in_f, gain=1.0):
|
|
w = self.randn(prefix + ".weight", (out_f, in_f), gain * (1.0 / in_f) ** 0.5).bfloat16()
|
|
q, s = quant_fp8_block(w)
|
|
self.add(prefix + ".weight", q)
|
|
self.add(prefix + ".scale", s)
|
|
|
|
def fp4(self, prefix, out_f, in_f, gain=1.0):
|
|
w = self.randn(prefix + ".weight", (out_f, in_f), gain * (1.0 / in_f) ** 0.5).bfloat16()
|
|
q, s = quant_fp4_row(w)
|
|
self.add(prefix + ".weight", q)
|
|
self.add(prefix + ".scale", s)
|
|
|
|
def f32(self, name, t):
|
|
self.add(name, t.float())
|
|
|
|
|
|
def build_tensors(cfg: dict, seed: int, norm_noise: float, attn_out_scale: float, ffn_out_scale: float) -> "OrderedDict[str, torch.Tensor]":
|
|
t = cfg["text_config"]
|
|
v = cfg["vision_config"]
|
|
f = Factory(seed, norm_noise)
|
|
dim, H, hd = t["hidden_size"], t["num_attention_heads"], t["head_dim"]
|
|
qr, olr, og = t["q_lora_rank"], t["o_lora_rank"], t["o_groups"]
|
|
inter, hc = t["moe_intermediate_size"], t["hc_mult"]
|
|
n_layer, nextn = t["num_hidden_layers"], t["num_nextn_predict_layers"]
|
|
ratios = t["compress_ratios"]
|
|
kv_src, idx_src = set(t["kv_source_layer_ids"]), set(t["index_source_layer_ids"])
|
|
ih, ihd = t["index_n_heads"], t["index_head_dim"]
|
|
vocab = t["vocab_size"]
|
|
mix_hc = (2 + hc) * hc
|
|
|
|
# ---- vision tower + aligner (BF16), real shapes ----
|
|
vd, vinter, ps = v["hidden_size"], v["intermediate_size"], v["patch_size"]
|
|
f.bf16("vision.patch_embed.proj.weight", vd, 3 * ps * ps)
|
|
f.bias("vision.patch_embed.proj.bias", vd)
|
|
for i in range(v["num_hidden_layers"]):
|
|
p = f"vision.blocks.{i}."
|
|
f.norm(p + "norm1.weight", vd)
|
|
f.bf16(p + "attn.wqkv.weight", 3 * vd, vd)
|
|
f.bias(p + "attn.wqkv.bias", 3 * vd)
|
|
f.bf16(p + "attn.wo.weight", vd, vd)
|
|
f.bias(p + "attn.wo.bias", vd)
|
|
f.norm(p + "norm2.weight", vd)
|
|
f.bf16(p + "mlp.w1.weight", 2 * vinter, vd)
|
|
f.bf16(p + "mlp.w2.weight", vd, vinter)
|
|
f.norm("vision.norm.weight", vd)
|
|
ds = v["downsample_ratio"]
|
|
f.bf16("aligner.w1.weight", dim, vd * ds * ds)
|
|
f.bias("aligner.w1.bias", dim)
|
|
f.bf16("aligner.w2.weight", dim, dim)
|
|
f.bias("aligner.w2.bias", dim)
|
|
|
|
# ---- global ----
|
|
f.bf16("embed.weight", vocab, dim, std=1.0)
|
|
for n in ("image_start", "image_end", "image_newline"):
|
|
f.add(n, f.randn(n, (dim,), 1.0).bfloat16())
|
|
f.norm("norm.weight", dim)
|
|
f.bf16("head.weight", vocab, dim)
|
|
|
|
def hc_params(p):
|
|
for which in ("attn", "ffn"):
|
|
f.f32(p + f"hc_{which}_fn", f.randn(p + f"hc_{which}_fn", (mix_hc, hc * dim), (1.0 / (hc * dim)) ** 0.5))
|
|
f.f32(p + f"hc_{which}_scale", torch.full((3,), 0.3) + f.randn(p + f"hc_{which}_scale", (3,), 0.05))
|
|
base = f.randn(p + f"hc_{which}_base", (mix_hc,), 0.5)
|
|
base[2 * hc :] += 2.0 * torch.eye(hc).flatten() # comb biased towards identity
|
|
f.f32(p + f"hc_{which}_base", base)
|
|
|
|
def attn_params(p, layer_id):
|
|
f.f32(p + "attn.attn_sink", f.randn(p + "attn.attn_sink", (H,), 0.1))
|
|
f.fp8(p + "attn.wq_a", qr, dim)
|
|
f.norm(p + "attn.q_norm.weight", qr)
|
|
f.fp8(p + "attn.wq_b", H * hd, qr)
|
|
f.fp8(p + "attn.wkv", hd, dim)
|
|
f.norm(p + "attn.kv_norm.weight", hd)
|
|
f.fp8(p + "attn.wo_a", og * olr, H * hd // og)
|
|
f.fp8(p + "attn.wo_b", dim, og * olr, gain=attn_out_scale)
|
|
backbone = layer_id is not None
|
|
if backbone and layer_id in kv_src:
|
|
f.bf16(p + "attn.compressor.wkv.weight", hd, dim)
|
|
if ratios[layer_id] > 1:
|
|
f.bf16(p + "attn.compressor.wgate.weight", hd, dim)
|
|
f.norm(p + "attn.compressor.norm.weight", hd)
|
|
if backbone and layer_id in idx_src:
|
|
f.fp8(p + "attn.indexer.wq_b", ih * ihd, qr)
|
|
f.bf16(p + "attn.indexer.weights_proj.weight", ih, dim)
|
|
if layer_id in kv_src: # only index-K owners carry wk / k_norm
|
|
f.bf16(p + "attn.indexer.wk.weight", ihd, hd)
|
|
f.norm(p + "attn.indexer.k_norm.weight", ihd)
|
|
|
|
def ffn_params(p, n_exp):
|
|
f.norm(p + "attn_norm.weight", dim)
|
|
f.norm(p + "ffn_norm.weight", dim)
|
|
f.bf16(p + "ffn.gate.weight", n_exp, dim)
|
|
f.bias(p + "ffn.gate.bias", n_exp, std=0.02, dtype=torch.float32)
|
|
f.bias(p + "ffn.gate.bias_vl", n_exp, std=0.02, dtype=torch.float32)
|
|
f.fp8(p + "ffn.shared_experts.w1", inter, dim)
|
|
f.fp8(p + "ffn.shared_experts.w3", inter, dim)
|
|
f.fp8(p + "ffn.shared_experts.w2", dim, inter, gain=ffn_out_scale)
|
|
for e in range(n_exp):
|
|
q = p + f"ffn.experts.{e}."
|
|
f.fp4(q + "w1", inter, dim)
|
|
f.fp4(q + "w3", inter, dim)
|
|
f.fp4(q + "w2", dim, inter, gain=ffn_out_scale)
|
|
|
|
primes = engram_primes(t["engram_layer_ids"], t["engram_max_ngram_size"], t["engram_n_heads"], t["engram_vocab_size"])
|
|
n_hash_cols = (t["engram_max_ngram_size"] - 1) * t["engram_n_heads"]
|
|
ehd = t["engram_head_dim"]
|
|
|
|
for L in range(n_layer):
|
|
p = f"layers.{L}."
|
|
hc_params(p)
|
|
attn_params(p, L)
|
|
ffn_params(p, t["n_routed_experts"])
|
|
if L in t["engram_layer_ids"]:
|
|
li = t["engram_layer_ids"].index(L)
|
|
rows = sum(sum(x) for x in primes[li])
|
|
assert rows == t["engram_num_embeddings"][li]
|
|
w = f.randn(p + "engram.embed.weight", (rows, ehd), 1.0).bfloat16()
|
|
q, s = quant_fp8_row(w)
|
|
f.add(p + "engram.embed.weight", q)
|
|
f.add(p + "engram.embed.scale", s)
|
|
f.add(p + "engram.q_weight", (1.0 + f.randn(p + "engram.q_weight", (hc, dim), 0.1)).bfloat16())
|
|
f.add(p + "engram.k_weight", (1.0 + f.randn(p + "engram.k_weight", (hc, dim), 0.1)).bfloat16())
|
|
# key half is RMS-normalised inside the gate, so the gain only shrinks the value written to the stream
|
|
f.fp8(p + "engram.wkv", dim * (hc + 1), n_hash_cols * ehd, gain=ffn_out_scale)
|
|
|
|
for i in range(nextn):
|
|
p = f"mtp.{i}."
|
|
hc_params(p)
|
|
attn_params(p, None)
|
|
ffn_params(p, t["dspark_n_routed_experts"])
|
|
if i == 0:
|
|
f.fp8(p + "main_proj", dim, dim * len(t["dspark_target_layer_ids"]))
|
|
f.norm(p + "main_norm.weight", dim)
|
|
if i == nextn - 1:
|
|
r = t["dspark_markov_rank"]
|
|
f.norm(p + "norm.weight", dim)
|
|
f.bf16(p + "markov_head.embed.weight", vocab, r, std=1.0)
|
|
f.bf16(p + "markov_head.head.weight", vocab, r)
|
|
f.bf16(p + "confidence_head.proj.weight", 1, dim + r)
|
|
return f.tensors
|
|
|
|
|
|
def make_config(real: dict) -> dict:
|
|
cfg = copy.deepcopy(real)
|
|
t = cfg["text_config"]
|
|
for k, val in TINY_TEXT_OVERRIDES.items():
|
|
assert k in t, k
|
|
t[k] = val
|
|
for k, val in TINY_VISION_OVERRIDES.items():
|
|
cfg["vision_config"][k] = val
|
|
assert len(t["compress_ratios"]) == t["num_hidden_layers"] + t["num_nextn_predict_layers"]
|
|
primes = engram_primes(t["engram_layer_ids"], t["engram_max_ngram_size"], t["engram_n_heads"], t["engram_vocab_size"])
|
|
t["engram_num_embeddings"] = [sum(sum(x) for x in per_layer) for per_layer in primes]
|
|
return cfg
|
|
|
|
|
|
def shard_of(name: str) -> int:
|
|
if name.startswith(("vision.", "aligner.")):
|
|
return 0
|
|
if name.startswith(("embed.", "head.", "norm.", "image_")):
|
|
return 1
|
|
if name.startswith("layers."):
|
|
return 2 if int(name.split(".")[1]) < 5 else 3
|
|
return 4 # mtp.*
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--out", required=True)
|
|
ap.add_argument("--tokenizer-dir", required=True)
|
|
ap.add_argument("--real-config", required=True, help="the released HF config.json")
|
|
ap.add_argument("--seed", type=int, default=20261020)
|
|
ap.add_argument("--norm-noise", type=float, default=0.0, help="std of N(1, s) norm weights (0 = all ones)")
|
|
# Gains on the projections that write into the residual stream. With plain N(0, 1/in) everywhere (1.0, 1.0)
|
|
# the reference is chaotic: FP8 activation quantization + top-k routing turn sub-ulp differences (e.g. another
|
|
# CPU thread count) into 20-50% logit differences. Small residual updates mimic trained models; the MoE/engram
|
|
# branch (discrete routing) is damped most, attention less so that the logits still depend on the context.
|
|
# Measured (threads 16 vs 7, teacher forced): logits rel err <= 0.04 while short/medium/long contexts ending in
|
|
# the same token differ by 0.24-0.49.
|
|
ap.add_argument("--attn-out-scale", type=float, default=0.5, help="gain of attn.wo_b")
|
|
ap.add_argument("--ffn-out-scale", type=float, default=1.0 / 16, help="gain of shared/routed expert w2 and engram wkv")
|
|
args = ap.parse_args()
|
|
torch.set_num_threads(min(16, os.cpu_count() or 1))
|
|
|
|
real = json.load(open(args.real_config))
|
|
cfg = make_config(real)
|
|
os.makedirs(args.out, exist_ok=True)
|
|
tensors = build_tensors(cfg, args.seed, args.norm_noise, args.attn_out_scale, args.ffn_out_scale)
|
|
|
|
# self-check of the FP4 packing against convert.py's decoding convention
|
|
for name, ten in tensors.items():
|
|
if name.endswith("experts.0.w1.weight") and name.startswith("layers.0."):
|
|
deq = dequant_fp4_row(ten, tensors[name.replace("weight", "scale")])
|
|
ref = torch.randn((ten.shape[0], ten.shape[1] * 2), generator=Factory(args.seed, 0).gen(name)) * (1.0 / (ten.shape[1] * 2)) ** 0.5
|
|
rel = (deq - ref.bfloat16().float()).norm() / ref.norm()
|
|
print(f"fp4 self-check {name}: rel err {rel:.4f}")
|
|
assert rel < 0.2
|
|
|
|
n_shards = 5
|
|
shards = [OrderedDict() for _ in range(n_shards)]
|
|
for name, ten in tensors.items():
|
|
shards[shard_of(name)][name] = ten
|
|
weight_map, total = {}, 0
|
|
for i, sd in enumerate(shards):
|
|
fname = f"model-{i + 1:05d}-of-{n_shards:05d}.safetensors"
|
|
save_file(sd, os.path.join(args.out, fname), metadata={"format": "pt"})
|
|
for name, ten in sd.items():
|
|
weight_map[name] = fname
|
|
total += ten.numel() * ten.element_size()
|
|
with open(os.path.join(args.out, "model.safetensors.index.json"), "w") as fh:
|
|
json.dump({"metadata": {"total_size": total}, "weight_map": dict(sorted(weight_map.items()))}, fh, indent=2)
|
|
with open(os.path.join(args.out, "config.json"), "w") as fh:
|
|
json.dump(cfg, fh, indent=2)
|
|
with open(os.path.join(args.out, "tiny_checkpoint.json"), "w") as fh:
|
|
json.dump({"generator": "examples/neuron/make_tiny_checkpoint.py", "seed": args.seed, "norm_noise": args.norm_noise,
|
|
"attn_out_scale": args.attn_out_scale, "ffn_out_scale": args.ffn_out_scale, "n_tensors": len(tensors), "total_bytes": total}, fh, indent=2)
|
|
for fn in ("tokenizer.json", "tokenizer_config.json"):
|
|
shutil.copyfile(os.path.join(args.tokenizer_dir, fn), os.path.join(args.out, fn))
|
|
print(f"wrote {len(tensors)} tensors, {total / 2**30:.2f} GiB to {args.out}")
|
|
print("engram_num_embeddings", cfg["text_config"]["engram_num_embeddings"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|