File size: 9,110 Bytes
e1acc15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | """Streaming loader: official 1017 tensors -> each rank's final native shards on its own device.
Nothing is assembled whole on a GPU: every expert matrix / dense tensor is read from the memory-mapped official shard,
cut on the host to the rank's block-aligned window, copied once to that rank's device and (expanded arms) decoded there
straight into its slot of the preallocated bank. Peak extra device memory while loading = one matrix being decoded
(~0.1 GB); peak extra host memory = one matrix (~6 MB) plus the page cache of the mapped shards.
Verification (fail loudly, never serve unverified bytes): every official tensor is re-hashed against
NATIVE_MANIFEST.json before use; packed shards are compared byte for byte with the host slice after the copy;
expanded shards are compared with formats.decode_fp4 on the host for a sample of experts per stage (all three matrices,
both ranks) and checked finite everywhere. FP8 dense scales become fp32 values (exact powers of two).
"""
from __future__ import annotations
import time
from dataclasses import dataclass, field
import numpy as np
import torch
from . import formats as F
from .official import OfficialShards
from .ops import ExpertBank, NativeDense
N_EXPERTS = 128
HIDDEN = 5120
INTER = 2304
@dataclass
class LoadStats:
seconds: float = 0.0
bytes_read: int = 0
tensors: int = 0
verified_full: int = 0
verified_sample: int = 0
peak_alloc: dict = field(default_factory=dict) # device -> torch max_memory_allocated during the load
peak_reserved: dict = field(default_factory=dict)
def _np(sh: OfficialShards, name: str, stats: LoadStats, verify: bool) -> np.ndarray:
if verify:
sh.verify(name)
info = sh.info(name)
raw = sh.raw(name)
stats.bytes_read += raw.nbytes
stats.tensors += 1
shape = info["shape"]
if info["dtype"] in ("I8", "F8_E4M3", "F8_E8M0"):
return np.asarray(raw).view(np.uint8).reshape(shape)
if info["dtype"] == "BF16":
return np.asarray(raw).view(np.uint16).reshape(shape)
if info["dtype"] == "F32":
return np.asarray(raw).view(np.float32).reshape(shape)
raise ValueError(f"{name}: dtype {info['dtype']}")
def _to(dev: torch.device, a: np.ndarray, dtype: torch.dtype | None = None) -> torch.Tensor:
t = torch.from_numpy(np.ascontiguousarray(a))
if dtype is torch.bfloat16:
t = t.view(torch.bfloat16)
return t.to(dev)
def _check_scales(name: str, s: np.ndarray) -> None:
if (s == 255).any() or (s == 0).any():
raise RuntimeError(f"{name}: E8M0 scale bytes 0/255 present (not served: the kernels assume 1..254)")
def load_expert_banks(sh: OfficialShards, stage: int, ranks: list[tuple[int, torch.device]], tp: int, kind: str,
stats: LoadStats, verify_sample: int = 4, n_experts: int = N_EXPERTS) -> dict[int, ExpertBank]:
"""Routed experts of mtp.<stage> for each (rank, device): rank r holds intermediate channels
[r*I/tp, (r+1)*I/tp) of w1/w3 (output rows) and of w2 (inputs, = packed bytes and 32-groups of the same window)."""
I_r = INTER // tp
if I_r % 64:
raise ValueError(f"tp {tp}: {I_r} channels per rank not a multiple of 64")
banks = {r: ExpertBank(f"mtp.{stage}.experts[r{r}]", kind, n_experts, HIDDEN, I_r, dev) for r, dev in ranks}
sample = set(np.linspace(0, n_experts - 1, verify_sample).astype(int).tolist()) if verify_sample else set()
for e in range(n_experts):
p = f"mtp.{stage}.ffn.experts.{e}."
w1, s1 = _np(sh, p + "w1.weight", stats, True), _np(sh, p + "w1.scale", stats, True)
w3, s3 = _np(sh, p + "w3.weight", stats, True), _np(sh, p + "w3.scale", stats, True)
w2, s2 = _np(sh, p + "w2.weight", stats, True), _np(sh, p + "w2.scale", stats, True)
for nm, s in (("w1", s1), ("w3", s3), ("w2", s2)):
_check_scales(p + nm, s)
for r, dev in ranks:
bk = banks[r]
rows = slice(r * I_r, (r + 1) * I_r)
pk = slice(r * I_r // 2, (r + 1) * I_r // 2)
sg = slice(r * I_r // 32, (r + 1) * I_r // 32)
parts = {"w1": (w1[rows], s1[rows]), "w3": (w3[rows], s3[rows]), "w2": (w2[:, pk], s2[:, sg])}
for j, nm in enumerate(("w1", "w3")):
wd, sd = _to(dev, parts[nm][0]), _to(dev, parts[nm][1])
if kind == "packed":
bk.w13[e, j].copy_(wd)
bk.s13[e, j].copy_(sd)
else:
bk.w13[e, j].copy_(F.decode_fp4(wd, sd).to(torch.bfloat16))
wd, sd = _to(dev, parts["w2"][0]), _to(dev, parts["w2"][1])
if kind == "packed":
bk.w2[e].copy_(wd)
bk.s2[e].copy_(sd)
else:
bk.w2[e].copy_(F.decode_fp4(wd, sd).to(torch.bfloat16))
# verification
if kind == "packed":
for got, want in ((bk.w13[e, 0], parts["w1"][0]), (bk.s13[e, 0], parts["w1"][1]),
(bk.w13[e, 1], parts["w3"][0]), (bk.s13[e, 1], parts["w3"][1]),
(bk.w2[e], parts["w2"][0]), (bk.s2[e], parts["w2"][1])):
if not np.array_equal(got.cpu().numpy(), np.ascontiguousarray(want)):
raise RuntimeError(f"{p}: rank {r} packed shard differs from the official bytes")
stats.verified_full += 6
elif e in sample:
for got, (wa, sa) in ((bk.w13[e, 0], parts["w1"]), (bk.w13[e, 1], parts["w3"]), (bk.w2[e], parts["w2"])):
want = F.decode_fp4(torch.from_numpy(np.ascontiguousarray(wa)),
torch.from_numpy(np.ascontiguousarray(sa)))
if not torch.equal(got.cpu().float(), want):
raise RuntimeError(f"{p}: rank {r} expanded shard is not the exact decode")
stats.verified_sample += 3
for r, dev in ranks:
if kind == "expanded": # one expert at a time: a whole-bank isfinite needs ~4 GiB of temporaries
for e in range(n_experts):
for t in (banks[r].w13[e], banks[r].w2[e]):
if not bool(torch.isfinite(t).all()):
raise RuntimeError(f"mtp.{stage} expert {e} rank {r}: non-finite expanded weights")
return banks
def load_dense(sh: OfficialShards, name: str, dev: torch.device, kind: str, stats: LoadStats,
rows: tuple[int, int] | None = None, cols: tuple[int, int] | None = None) -> NativeDense:
"""FP8 + E8M0 32x32 tensor ``name`` (``.weight`` / ``.scale``), the block-aligned window rows x cols, as ``fp8``
(bytes + fp32 scale values) or ``bf16`` (exact decode). Verified: bytes / exact decode against the host."""
w = _np(sh, name + ".weight", stats, True)
s = _np(sh, name + ".scale", stats, True)
_check_scales(name, s)
N, K = w.shape
r0, r1 = rows or (0, N)
c0, c1 = cols or (0, K)
if r0 % 32 or c0 % 32 or (r1 - r0) % 32 or (c1 - c0) % 32:
raise ValueError(f"{name}: window rows {r0}:{r1} cols {c0}:{c1} not on 32x32 scale blocks")
ws = np.ascontiguousarray(w[r0:r1, c0:c1])
ss = np.ascontiguousarray(s[r0 // 32:r1 // 32, c0 // 32:c1 // 32])
tag = f"{name}[{r0}:{r1},{c0}:{c1}]"
if kind == "fp8":
if ((ws == 0x7F) | (ws == 0xFF)).any():
raise RuntimeError(f"{name}: NaN e4m3 codes")
sv = F.e8m0_to_f32(torch.from_numpy(ss))
d = NativeDense(tag, "fp8", _to(dev, ws), sv.to(dev))
if not np.array_equal(d.w.cpu().numpy(), ws):
raise RuntimeError(f"{tag}: device bytes differ from the official bytes")
else:
full = F.decode_fp8_block(torch.from_numpy(ws), torch.from_numpy(ss))
d = NativeDense(tag, "bf16", full.to(torch.bfloat16).to(dev), None)
if not torch.equal(d.w.cpu().float(), full):
raise RuntimeError(f"{tag}: expanded weights are not the exact decode")
stats.verified_full += 1
return d
def load_bf16(sh: OfficialShards, name: str, dev: torch.device, stats: LoadStats,
rows: tuple[int, int] | None = None) -> torch.Tensor:
"""An official BF16 tensor (e.g. head.weight), optionally a row window, verified by sha256 of the whole tensor."""
a = _np(sh, name, stats, True)
if rows:
a = a[rows[0]:rows[1]]
t = _to(dev, a, torch.bfloat16)
stats.verified_full += 1
return t
def track_peaks(devs: list[torch.device]):
"""Context helper: reset and later read torch's peak counters on each device."""
class _P:
def __enter__(self):
self.t0 = time.time()
for d in devs:
torch.cuda.synchronize(d)
torch.cuda.reset_peak_memory_stats(d)
return self
def __exit__(self, *a):
self.seconds = time.time() - self.t0
self.alloc = {str(d): torch.cuda.max_memory_allocated(d) for d in devs}
self.reserved = {str(d): torch.cuda.max_memory_reserved(d) for d in devs}
return False
return _P()
|