Frank
8343397584
The reader opens every index-named file in sorted order and emplaces EVERY tensor each file physically holds, first-wins (checkpoint.cpp add_shard + shard_list). The raw model-000NN shards carry tensors named ...ple.ple_embedding.ngram_embedding.shard_0.weight ([2500012,160] fp8) - the identical name the plugin asks the merged table under - so ple.merged.safetensors lost the emplace race and the plan read one raw shard as the codes (declared [320001536,160] vs codes [2500012,160]). Rename to 00-ple.merged.safetensors: '00-' wins the sort, the merged bf16 tensor wins the name, table_bf16 declares and the rtn rule makes the fp8 container form.
557 lines
23 KiB
Python
Executable File
557 lines
23 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""stack-qwen38-official-fp8.py — reshape the official Qwen/Qwen3.8-Flash-Next-FP8
|
|
checkpoint into the form the shipped radiance qwen4exp plugin maps.
|
|
|
|
Run INSIDE the staging dir's tree (python3 + numpy on the host; data never
|
|
leaves the box). Two transforms, both VALUE-PRESERVING:
|
|
|
|
1. EXPERTS: the shipped plugin's stacked layout names ONE tensor per family
|
|
(`mlp.experts.gate_up_proj` [E, 2F, H], `mlp.experts.down_proj` [E, H, F])
|
|
plus a `<family>_scale_inv` plane at the same index (NO `.weight` in the
|
|
stacked name — scale_name() appends _scale_inv to the tensor name verbatim),
|
|
while the official FP8 repo ships per-expert unfused tensors
|
|
(`experts.<e>.{gate,up,down}_proj.weight` + scale). Stacking is pure byte
|
|
concatenation along dim 0 — gate rows
|
|
first inside each expert's gate_up slab, which is the order the block
|
|
wants (rad_block_moe_fp8.h: "a stacked gate_up is ALREADY FUSED in the
|
|
order this block wants"). No requantisation: bytes are moved, not decoded.
|
|
|
|
2. PLE TABLE: the plugin's fp8 branch wants the table's codes beside a
|
|
ONE-ELEMENT bf16 scale plane (`blk.%d.ple_ngram.scale`, view {1}); a
|
|
per-shard-blocked scale can never fit that view, and `rtn` re-encoding the
|
|
RAW scaleless codes would overflow (v / 1.99e-4 exceeds e4m3's 448 max).
|
|
So staging decodes the scaleless shards with the checkpoint's OWN bf16
|
|
scale into a bf16 table (exact: every e4m3 code value and the scale are
|
|
exactly representable in bf16 to within the code's own 2^-3 grid), and the
|
|
convert's rtn rule re-encodes at that same scale — the nearest-code
|
|
distance is 2^-3 while the bf16 rounding perturbs by 2^-9, so every
|
|
original code is restored exactly and the container gets the checkpoint's
|
|
own codes plus the one-element scale plane.
|
|
|
|
The convert then runs with NO recipe except:
|
|
blk.*.ple_ngram.weight=rtn:codes=fp8_e4m3,block=*x*,scale=bf16,rule=fixed,scale_value=<s>
|
|
with <s> read from the checkpoint and written to ple_scale.txt.
|
|
|
|
Idempotent: re-runs skip when the staged index already carries the stacked
|
|
names, the merged bf16 shard_0, and split_ngram_parts == 1. --force redoes.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import mmap
|
|
import multiprocessing as mp
|
|
import os
|
|
import re
|
|
import struct
|
|
import sys
|
|
import time
|
|
import numpy as np
|
|
F8, BF16 = "F8_E4M3", "BF16"
|
|
|
|
|
|
def log(msg):
|
|
print(msg, flush=True)
|
|
|
|
|
|
# ---------------------------------------------------------------- safetensors
|
|
|
|
def parse_header(path):
|
|
"""Return (header_len, {name: (dtype, shape, data_offsets)})."""
|
|
with open(path, "rb") as fh:
|
|
(hlen,) = struct.unpack("<Q", fh.read(8))
|
|
hdr = json.loads(fh.read(hlen))
|
|
return hlen, hdr
|
|
|
|
|
|
def write_safetensors(path, tensors, data_iter):
|
|
"""tensors: [(name, dtype, shape, nbytes)]; data_iter yields bytes per tensor."""
|
|
hdr = {}
|
|
off = 0
|
|
for name, dtype, shape, nbytes in tensors:
|
|
hdr[name] = {"dtype": dtype, "shape": list(shape),
|
|
"data_offsets": [off, off + nbytes]}
|
|
off += nbytes
|
|
blob = json.dumps(hdr, separators=(",", ":")).encode("utf-8")
|
|
pad = (-len(blob)) % 8 # keep the data start 8-byte aligned
|
|
blob += b" " * pad
|
|
with open(path, "wb") as fh:
|
|
fh.write(struct.pack("<Q", len(blob)))
|
|
fh.write(blob)
|
|
for _, _, _, nbytes in tensors:
|
|
fh.write(next(data_iter).tobytes())
|
|
fh.flush()
|
|
os.fsync(fh.fileno())
|
|
return 8 + len(blob) + sum(t[3] for t in tensors)
|
|
|
|
|
|
# ------------------------------------------------------------------- preflight
|
|
|
|
def build_meta(stage, index):
|
|
"""name -> (file, global_byte_offset, nbytes, dtype, shape) for every tensor."""
|
|
files = sorted(set(index["weight_map"].values()))
|
|
meta = {}
|
|
for f in files:
|
|
path = os.path.join(stage, f)
|
|
hlen, hdr = parse_header(path)
|
|
for name, e in hdr.items():
|
|
if name == "__metadata__":
|
|
continue
|
|
b0, b1 = e["data_offsets"]
|
|
meta[name] = (f, 8 + hlen + b0, b1 - b0, e["dtype"], e["shape"])
|
|
return meta
|
|
|
|
|
|
def e4m3_lut():
|
|
lut = np.zeros(256, np.float32)
|
|
for b in range(256):
|
|
sign = -1.0 if b & 0x80 else 1.0
|
|
e = (b >> 3) & 0xF
|
|
m = b & 0x7
|
|
if e == 15 and m == 7:
|
|
v = np.float32(np.nan) # OCP e4m3fn: only 0x7F/0xFF are NaN
|
|
elif e == 0:
|
|
v = sign * (m / 8.0) * 2.0 ** -6
|
|
else:
|
|
v = sign * (1.0 + m / 8.0) * 2.0 ** (e - 7)
|
|
lut[b] = v
|
|
assert lut[0x38] == 1.0 and lut[0x7E] == 448.0
|
|
assert np.isnan(lut[0x7F]) and lut[0x00] == 0.0
|
|
assert lut[0x08] == 2.0 ** -6 and lut[0x80] == 0.0
|
|
return lut
|
|
|
|
|
|
def f32_to_bf16_rne(v):
|
|
"""float32 array -> uint16 bf16 bits, round-to-nearest-even."""
|
|
q = np.ascontiguousarray(v, dtype=np.float32)
|
|
q[~np.isfinite(q)] = 0.0
|
|
bits = q.view(np.uint32)
|
|
r = bits + 0x7FFF + ((bits >> 16) & 1)
|
|
return (r >> 16).astype(np.uint16)
|
|
|
|
|
|
# ------------------------------------------------------------------- workers
|
|
|
|
def stack_prefix(job):
|
|
"""Byte-stack one layer's 512 experts into a 4-tensor safetensors file."""
|
|
(stage, pre, idx, n_e, shapes, meta) = job
|
|
(f2, h) = shapes["ffn2_h"] # 2*ffn, n_embd
|
|
(f,) = shapes["ffn"] # ffn
|
|
gs, ds = shapes["gs"], shapes["ds"] # scale shapes [5,20] and [20,5]
|
|
|
|
# byte widths: fp8 is 1 B/element, bf16 2 B; the fused gate_up scale is
|
|
# the gate and up scale planes concatenated (gate rows first)
|
|
g_bytes = (f2 // 2) * h
|
|
gs_bytes = gs[0] * gs[1] * 2 # ONE projection's scale plane
|
|
gus_bytes = 2 * gs_bytes # fused scale per expert
|
|
dn_bytes = h * f
|
|
ds_bytes = ds[0] * ds[1] * 2
|
|
|
|
gu = np.empty(n_e * f2 * h, np.uint8)
|
|
gus = np.empty(n_e * gus_bytes, np.uint8)
|
|
dn = np.empty(n_e * dn_bytes, np.uint8)
|
|
dns = np.empty(n_e * ds_bytes, np.uint8)
|
|
|
|
mmaps = {}
|
|
|
|
def view(name):
|
|
f_, off, ln, dt, shp = meta[name]
|
|
if f_ not in mmaps:
|
|
fh = open(os.path.join(stage, f_), "rb")
|
|
mmaps[f_] = mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ)
|
|
return memoryview(mmaps[f_])[off:off + ln], (dt, shp)
|
|
|
|
def need(name, dt, shp):
|
|
mv, (got_dt, got_shp) = view(name)
|
|
assert got_dt == dt and list(got_shp) == list(shp), f"{name}: {got_dt}{got_shp}"
|
|
return mv
|
|
|
|
for e in range(n_e):
|
|
base = f"{pre}.mlp.experts.{e}"
|
|
slab = gu[e * f2 * h:(e + 1) * f2 * h]
|
|
slab[:g_bytes] = need(f"{base}.gate_proj.weight", F8, [f2 // 2, h])
|
|
slab[g_bytes:] = need(f"{base}.up_proj.weight", F8, [f2 // 2, h])
|
|
sslab = gus[e * gus_bytes:(e + 1) * gus_bytes]
|
|
sslab[:gs_bytes] = need(f"{base}.gate_proj.weight_scale_inv", BF16, gs)
|
|
sslab[gs_bytes:] = need(f"{base}.up_proj.weight_scale_inv", BF16, gs)
|
|
dn[e * dn_bytes:(e + 1) * dn_bytes] = need(f"{base}.down_proj.weight", F8, [h, f])
|
|
dns[e * ds_bytes:(e + 1) * ds_bytes] = need(f"{base}.down_proj.weight_scale_inv", BF16, ds)
|
|
|
|
for mm in mmaps.values():
|
|
mm.close()
|
|
|
|
# byte-identity check on two sampled experts (first and last), all four
|
|
# tensors — the stack IS the checkpoint bytes, so they must compare equal
|
|
for e in (0, n_e - 1):
|
|
base = f"{pre}.mlp.experts.{e}"
|
|
s = e * f2 * h
|
|
assert bytes(gu[s:s + g_bytes]) == bytes(stack_src(f"{base}.gate_proj.weight", meta, stage))
|
|
assert bytes(gu[s + g_bytes:s + f2 * h]) == bytes(stack_src(f"{base}.up_proj.weight", meta, stage))
|
|
t = e * gus_bytes
|
|
assert bytes(gus[t:t + gs_bytes]) == bytes(stack_src(f"{base}.gate_proj.weight_scale_inv", meta, stage))
|
|
assert bytes(gus[t + gs_bytes:t + gus_bytes]) == bytes(stack_src(f"{base}.up_proj.weight_scale_inv", meta, stage))
|
|
d = e * dn_bytes
|
|
assert bytes(dn[d:d + dn_bytes]) == bytes(stack_src(f"{base}.down_proj.weight", meta, stage))
|
|
u = e * ds_bytes
|
|
assert bytes(dns[u:u + ds_bytes]) == bytes(stack_src(f"{base}.down_proj.weight_scale_inv", meta, stage))
|
|
|
|
out = os.path.join(stage, f"stacked.{idx:03d}.safetensors")
|
|
tensors = [
|
|
(f"{pre}.mlp.experts.gate_up_proj", F8, [n_e, f2, h], n_e * f2 * h),
|
|
(f"{pre}.mlp.experts.gate_up_proj_scale_inv", BF16,
|
|
[n_e, gs[0] * 2, gs[1]], n_e * gus_bytes),
|
|
(f"{pre}.mlp.experts.down_proj", F8, [n_e, h, f], n_e * h * f),
|
|
(f"{pre}.mlp.experts.down_proj_scale_inv", BF16,
|
|
[n_e, ds[0], ds[1]], n_e * ds_bytes),
|
|
]
|
|
bufs = iter([gu, gus, dn, dns])
|
|
written = write_safetensors(out, tensors, bufs)
|
|
return (pre, written)
|
|
|
|
|
|
def stack_src(name, meta, stage):
|
|
"""Re-open a source tensor for the sampled-expert check."""
|
|
f_, off, ln, dt, shp = meta[name]
|
|
fh = open(os.path.join(stage, f_), "rb")
|
|
mm = mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ)
|
|
fh.close() # mmap keeps the mapping alive
|
|
return mm[off:off + ln]
|
|
|
|
|
|
def fill_ple_shard(job):
|
|
"""Decode one fp8 shard with the table scale, write bf16 rows into the mmap."""
|
|
(src_path, src_off, src_rows, dst_path, dst_off, scale_f32,
|
|
cols, chunk_rows) = job
|
|
lut = e4m3_lut()
|
|
with open(src_path, "rb") as fh:
|
|
sm = mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ)
|
|
with open(dst_path, "r+b") as dfh:
|
|
dm = mmap.mmap(dfh.fileno(), 0, access=mmap.ACCESS_WRITE)
|
|
for r0 in range(0, src_rows, chunk_rows):
|
|
r1 = min(r0 + chunk_rows, src_rows)
|
|
codes = np.frombuffer(sm, dtype=np.uint8,
|
|
count=(r1 - r0) * cols,
|
|
offset=src_off + r0 * cols)
|
|
q = lut.take(codes)
|
|
del codes # release the mmap export before sm.close()
|
|
q[~np.isfinite(q)] = 0.0
|
|
q *= np.float32(scale_f32)
|
|
packed = f32_to_bf16_rne(q)
|
|
dm[dst_off + r0 * cols * 2: dst_off + r1 * cols * 2] = packed.tobytes()
|
|
dm.flush()
|
|
os.fsync(dfh.fileno())
|
|
dm.close()
|
|
sm.close()
|
|
|
|
|
|
# ---------------------------------------------------------------------- main
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
|
|
ap.add_argument("--src", default=None,
|
|
help="raw checkpoint dir (documentation only; unused)")
|
|
ap.add_argument("--stage", required=True, help="staging dir to reshape in place")
|
|
ap.add_argument("--force", action="store_true", help="redo even if already staged")
|
|
ap.add_argument("--workers", type=int, default=min(8, os.cpu_count() or 1))
|
|
args = ap.parse_args()
|
|
stage = args.stage
|
|
t0 = time.time()
|
|
|
|
index_path = os.path.join(stage, "model.safetensors.index.json")
|
|
cfg_path = os.path.join(stage, "config.json")
|
|
if not (os.path.isfile(index_path) and os.path.isfile(cfg_path)):
|
|
log(f"ERROR: {stage} is not a staged checkpoint (missing index/config)")
|
|
return 2
|
|
with open(index_path) as fh:
|
|
index = json.load(fh)
|
|
with open(cfg_path) as fh:
|
|
cfg = json.load(fh)
|
|
wm = index["weight_map"]
|
|
|
|
meta = build_meta(stage, index)
|
|
|
|
# ---- collect prefixes -------------------------------------------------
|
|
pre_re = re.compile(r"^(?P<pre>.+)\.mlp\.experts\.(?P<e>\d+)\.gate_proj\.weight$")
|
|
pres, n_e = {}, 0
|
|
for k in wm:
|
|
m = pre_re.match(k)
|
|
if m:
|
|
pres.setdefault(m.group("pre"), set()).add(int(m.group("e")))
|
|
n_e = max(n_e, int(m.group("e")) + 1)
|
|
prefixes = sorted(pres)
|
|
if not prefixes:
|
|
log("ERROR: no per-expert tensors found in the staged index")
|
|
return 2
|
|
mtp = [p for p in prefixes if p.startswith("mtp.")]
|
|
trunk = [p for p in prefixes if p.startswith("model.language_model.layers.")]
|
|
log(f"expert prefixes: {len(prefixes)} (trunk {len(trunk)}, mtp {len(mtp)}), "
|
|
f"{n_e} experts each")
|
|
|
|
# shapes from expert 0 of the first prefix
|
|
pre0 = prefixes[0]
|
|
shapes = {}
|
|
gdt, gsh = meta[f"{pre0}.mlp.experts.0.gate_proj.weight"][3:5]
|
|
ddt, dsh = meta[f"{pre0}.mlp.experts.0.down_proj.weight"][3:5]
|
|
gsd, gssh = meta[f"{pre0}.mlp.experts.0.gate_proj.weight_scale_inv"][3:5]
|
|
dsd, dssh = meta[f"{pre0}.mlp.experts.0.down_proj.weight_scale_inv"][3:5]
|
|
scale_law = (gdt == F8 and ddt == F8 and gsd == BF16 and dsd == BF16
|
|
and list(gssh) == [(gsh[0] + 127) // 128, (gsh[1] + 127) // 128]
|
|
and list(dssh) == [(dsh[0] + 127) // 128, (dsh[1] + 127) // 128])
|
|
if not scale_law:
|
|
log(f"ERROR: unexpected expert shapes: gate {gdt}{gsh} scale {gsd}{gssh} "
|
|
f"down {ddt}{dsh} scale {dsd}{dssh} (scale law 128x128)")
|
|
return 2
|
|
shapes = {"ffn2_h": (gsh[0] * 2, gsh[1]), "ffn": (dsh[1],),
|
|
"gs": tuple(gssh), "ds": tuple(dssh)}
|
|
f2, h = shapes["ffn2_h"]
|
|
f = shapes["ffn"][0]
|
|
log(f"shapes: gate_up [{n_e},{f2},{h}] F8, down [{n_e},{h},{f}] F8, "
|
|
f"scales {gssh}x2 / {dssh} bf16")
|
|
|
|
# ---- PLE table --------------------------------------------------------
|
|
ple_re = re.compile(r"^(?P<pre>.+)\.ple\.ple_embedding\.ngram_embedding\.shard_(?P<i>\d+)\.weight$")
|
|
ple_pre, shard_ids = None, {}
|
|
for k in wm:
|
|
m = ple_re.match(k)
|
|
if m:
|
|
ple_pre = m.group("pre")
|
|
shard_ids[int(m.group("i"))] = k
|
|
if ple_pre is None:
|
|
log("ERROR: no PLE table shards found")
|
|
return 2
|
|
n_shards = len(shard_ids)
|
|
if sorted(shard_ids) != list(range(n_shards)):
|
|
log(f"ERROR: PLE shard ids not contiguous 0..{n_shards - 1}")
|
|
return 2
|
|
sc_name = f"{ple_pre}.ple.ple_embedding.ngram_embedding.weight_scale"
|
|
sdt, ssh, sln = meta[sc_name][3], meta[sc_name][4], meta[sc_name][2]
|
|
if sdt != BF16 or list(ssh) != [1]:
|
|
log(f"ERROR: PLE scale is {sdt}{ssh}, expected BF16 [1]")
|
|
return 2
|
|
s_bits = struct.unpack("<H", bytes(view_src_tuple(meta, sc_name, stage)))[0]
|
|
scale_f32 = struct.unpack("<f", struct.pack("<I", s_bits << 16))[0]
|
|
log(f"PLE: {ple_pre}, {n_shards} shards, scale bf16 = {scale_f32!r}")
|
|
|
|
rows = [meta[shard_ids[i]][4][0] for i in range(n_shards)]
|
|
total_rows = sum(rows)
|
|
hd = meta[shard_ids[0]][4][1]
|
|
if any(meta[shard_ids[i]][4][1] != hd for i in range(n_shards)):
|
|
log("ERROR: PLE shards disagree on HD")
|
|
return 2
|
|
if hd % 128 != 0 and hd != 160:
|
|
log(f"WARNING: PLE HD={hd} is neither 160 nor a multiple of 128")
|
|
|
|
# ---- idempotence gate -------------------------------------------------
|
|
stacked_names = [f"{p}.mlp.experts.gate_up_proj" for p in prefixes]
|
|
merged_name = f"{ple_pre}.ple.ple_embedding.ngram_embedding.shard_0.weight"
|
|
already = (all(n in wm for n in stacked_names)
|
|
and merged_name in wm
|
|
and find_split(cfg) is not None and read_split(cfg) == 1)
|
|
if already and not args.force:
|
|
log("already staged (stacked names, merged shard_0 and split_ngram_parts=1 "
|
|
"all present) — nothing to do")
|
|
return 0
|
|
|
|
# ---- disk gate --------------------------------------------------------
|
|
stacked_bytes = len(prefixes) * n_e * (f2 * h + h * f)
|
|
ple_bytes = total_rows * hd * 2
|
|
need = stacked_bytes + ple_bytes
|
|
st = os.statvfs(stage)
|
|
avail = st.f_bavail * st.f_frsize
|
|
if avail < need * 1.05:
|
|
log(f"ERROR: {avail / 2**30:.1f} GiB free < {need / 2**30:.1f} GiB needed (+5%)")
|
|
return 2
|
|
log(f"disk: {avail / 2**30:.1f} GiB free, need {need / 2**30:.1f} GiB "
|
|
f"(experts {stacked_bytes / 2**30:.1f}, ple {ple_bytes / 2**30:.1f})")
|
|
|
|
# ---- stage experts ----------------------------------------------------
|
|
jobs = []
|
|
for idx, pre in enumerate(prefixes):
|
|
jmeta = {n: meta[n] for n in meta if n.startswith(f"{pre}.mlp.experts.")}
|
|
jobs.append((stage, pre, idx, n_e, shapes, jmeta))
|
|
ctx = mp.get_context("fork")
|
|
written_total = 0
|
|
with ctx.Pool(args.workers) as pool:
|
|
for pre, nbytes in pool.imap_unordered(stack_prefix, jobs):
|
|
written_total += nbytes
|
|
log(f" stacked {pre} ({nbytes / 2**30:.2f} GiB)")
|
|
log(f"experts stacked: {written_total / 2**30:.2f} GiB in {len(prefixes)} files")
|
|
|
|
# ---- stage PLE --------------------------------------------------------
|
|
# The file name sorts BEFORE the model-000NN shards ON PURPOSE: the reader
|
|
# opens every index-named file in sorted order and emplaces EVERY tensor
|
|
# it physically holds, first-wins (checkpoint.cpp add_shard). The raw
|
|
# shards carry tensors of the SAME name as the merged one
|
|
# (...ple.ple_embedding.ngram_embedding.shard_0.weight, [2500012,160]
|
|
# fp8), so ple.merged.safetensors lost the race and the plan read one raw
|
|
# shard as the codes. "00-" wins the sort and the name.
|
|
out_name = "00-ple.merged.safetensors"
|
|
out_path = os.path.join(stage, out_name)
|
|
data_bytes = total_rows * hd * 2
|
|
hdr = {merged_name: {"dtype": BF16, "shape": [total_rows, hd],
|
|
"data_offsets": [0, data_bytes]}}
|
|
blob = json.dumps(hdr, separators=(",", ":")).encode("utf-8")
|
|
pad = (-len(blob)) % 8
|
|
blob += b" " * pad
|
|
hlen = len(blob)
|
|
with open(out_path, "wb") as fh:
|
|
fh.write(struct.pack("<Q", hlen))
|
|
fh.write(blob)
|
|
fh.truncate(8 + hlen + data_bytes)
|
|
jobs = []
|
|
row0 = 0
|
|
for i in range(n_shards):
|
|
f_, off, ln, _, shp = meta[shard_ids[i]]
|
|
jobs.append((os.path.join(stage, f_), off, shp[0], out_path,
|
|
8 + hlen + row0 * hd * 2, scale_f32, hd, 131072))
|
|
row0 += shp[0]
|
|
assert row0 == total_rows
|
|
with ctx.Pool(args.workers) as pool:
|
|
pool.map(fill_ple_shard, jobs)
|
|
log(f"PLE merged: {data_bytes / 2**30:.2f} GiB bf16 "
|
|
f"({total_rows} rows x {hd} cols)")
|
|
|
|
# ---- post-verify ------------------------------------------------------
|
|
verify(stage, prefixes, n_e, shapes, meta, shard_ids, rows, hd,
|
|
scale_f32, out_path, hlen)
|
|
|
|
# ---- rewrite index + config -------------------------------------------
|
|
new_wm = {k: v for k, v in wm.items() if k not in set(shard_ids.values())}
|
|
for idx, pre in enumerate(prefixes):
|
|
fname = f"stacked.{idx:03d}.safetensors"
|
|
for t in (f"{pre}.mlp.experts.gate_up_proj",
|
|
f"{pre}.mlp.experts.gate_up_proj_scale_inv",
|
|
f"{pre}.mlp.experts.down_proj",
|
|
f"{pre}.mlp.experts.down_proj_scale_inv"):
|
|
new_wm[t] = fname
|
|
new_wm[merged_name] = out_name
|
|
total = sum(meta_size(meta, k) for k in new_wm
|
|
if k in meta) # originals
|
|
# new tensors' sizes
|
|
total += written_total + data_bytes
|
|
index["weight_map"] = new_wm
|
|
if "metadata" in index and isinstance(index["metadata"], dict):
|
|
index["metadata"]["total_size"] = total
|
|
# ATOMIC REPLACE, never open(path, "w"): the stage's index is a symlink
|
|
# to the RAW checkpoint (staging symlinks everything but config.json) —
|
|
# open("w") would rewrite the raw file through the link. os.replace
|
|
# swaps the link itself for a real file; the raw tree stays pristine.
|
|
for path, payload in ((index_path, json.dumps(index, indent=2)),
|
|
(cfg_path, json.dumps(cfg, indent=2))):
|
|
tmp = path + ".tmp"
|
|
with open(tmp, "w") as fh:
|
|
fh.write(payload)
|
|
os.replace(tmp, path)
|
|
n = set_split(cfg, 1)
|
|
log(f"config: split_ngram_parts=1 written at {n} location(s)")
|
|
|
|
with open(os.path.join(stage, "ple_scale.txt"), "w") as fh:
|
|
fh.write(repr(scale_f32) + "\n")
|
|
|
|
dt = time.time() - t0
|
|
log(f"staged in {dt / 60:.1f} min: {written_total / 2**30:.1f} GiB experts + "
|
|
f"{data_bytes / 2**30:.1f} GiB PLE; scale {scale_f32!r} -> ple_scale.txt")
|
|
return 0
|
|
|
|
|
|
def meta_size(meta, name):
|
|
_, _, ln, _, _ = meta[name]
|
|
return ln
|
|
|
|
|
|
def view_src_tuple(meta, name, stage):
|
|
f_, off, ln, _, _ = meta[name]
|
|
with open(os.path.join(stage, f_), "rb") as fh:
|
|
fh.seek(off)
|
|
return fh.read(ln)
|
|
|
|
|
|
def find_split(cfg):
|
|
"""Locate split_ngram_parts keys (recursive); return list of paths."""
|
|
hits = []
|
|
|
|
def walk(d, path):
|
|
if isinstance(d, dict):
|
|
for k, v in d.items():
|
|
if k == "split_ngram_parts":
|
|
hits.append(path + [k])
|
|
walk(v, path + [k])
|
|
walk(cfg, [])
|
|
return hits
|
|
|
|
|
|
def read_split(cfg):
|
|
hits = find_split(cfg)
|
|
if not hits:
|
|
return None
|
|
d = cfg
|
|
for k in hits[0][:-1]:
|
|
d = d[k]
|
|
return d[hits[0][-1]]
|
|
|
|
|
|
def set_split(cfg, value):
|
|
hits = find_split(cfg)
|
|
for path in hits:
|
|
d = cfg
|
|
for k in path[:-1]:
|
|
d = d[k]
|
|
d[path[-1]] = value
|
|
if not hits: # plugin's meta_int2 falls back to the top-level key
|
|
cfg["split_ngram_parts"] = value
|
|
hits.append(["split_ngram_parts"])
|
|
return hits
|
|
|
|
|
|
def verify(stage, prefixes, n_e, shapes, meta, shard_ids, rows, hd,
|
|
scale_f32, out_path, hlen):
|
|
"""Cheap strong checks: headers parse, sizes match, sampled bytes identical."""
|
|
f2, h = shapes["ffn2_h"]
|
|
f = shapes["ffn"][0]
|
|
gs, ds = shapes["gs"], shapes["ds"]
|
|
# headers of every new file
|
|
for idx, pre in enumerate(prefixes):
|
|
path = os.path.join(stage, f"stacked.{idx:03d}.safetensors")
|
|
hlen2, hdr = parse_header(path)
|
|
want = {
|
|
f"{pre}.mlp.experts.gate_up_proj": (F8, [n_e, f2, h]),
|
|
f"{pre}.mlp.experts.gate_up_proj_scale_inv":
|
|
(BF16, [n_e, gs[0] * 2, gs[1]]),
|
|
f"{pre}.mlp.experts.down_proj": (F8, [n_e, h, f]),
|
|
f"{pre}.mlp.experts.down_proj_scale_inv": (BF16, [n_e, ds[0], ds[1]]),
|
|
}
|
|
for name, (dt, shp) in want.items():
|
|
e = hdr[name]
|
|
assert e["dtype"] == dt and list(e["shape"]) == shp, name
|
|
assert e["data_offsets"][1] - e["data_offsets"][0] == \
|
|
int(np.prod(shp)) * (1 if dt == F8 else 2), name
|
|
_, phdr = parse_header(out_path)
|
|
e = phdr[next(iter(phdr))]
|
|
assert e["dtype"] == BF16 and list(e["shape"]) == [sum(rows), hd]
|
|
|
|
# sampled PLE rows: recompute and compare bytes exactly
|
|
lut = e4m3_lut()
|
|
out_off = 8 + hlen
|
|
for i in (0, len(shard_ids) // 2, len(shard_ids) - 1):
|
|
name = shard_ids[i]
|
|
f_, off, ln, _, shp = meta[name]
|
|
with open(os.path.join(stage, f_), "rb") as fh:
|
|
mm = mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ)
|
|
for r in (0, shp[0] - 1):
|
|
codes = np.frombuffer(mm, dtype=np.uint8, count=hd,
|
|
offset=off + r * hd)
|
|
q = lut.take(codes) * np.float32(scale_f32)
|
|
del codes # release the mmap export before mm.close()
|
|
want = f32_to_bf16_rne(q).tobytes()
|
|
got_row = sum(rows[:i]) + r
|
|
with open(out_path, "rb") as ofh:
|
|
ofh.seek(out_off + got_row * hd * 2)
|
|
got = ofh.read(hd * 2)
|
|
assert got == want, f"PLE row {got_row} mismatch"
|
|
mm.close()
|
|
log("verify: new-file headers, sizes and sampled PLE rows all exact; "
|
|
"expert slabs were byte-checked in-stack")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|