File size: 2,633 Bytes
0693474 | 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 | #!/usr/bin/env python3
"""Per-rank Engram row ranges, computed exactly as the engine does (common/engram.py: EngramLayout primes,
ParallelEngramEmbedding head sharding: part_n_hash_cols = cdiv(n_hash_cols, tp), head_start = rank * part,
vocab_start_idx = sum(head_sizes[:head_start]), part_num_embeddings = vocab_end - vocab_start).
python3 layout.py config.json [--tp 2] -> JSON {layer: [{rank, row_start, num_rows}, ...]}
No torch/vllm import: the prime search below is a copy of _is_prime / find_next_prime / EngramLayout.__init__.
"""
import argparse
import json
def _is_prime(n: int) -> bool:
if n < 2:
return False
for p in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37):
if n % p == 0:
return n == p
d, r = n - 1, 0
while d % 2 == 0:
d //= 2
r += 1
for a in (2, 7, 61):
x = pow(a, d, n)
if x in (1, n - 1):
continue
for _ in range(r - 1):
x = x * x % n
if x == n - 1:
break
else:
return False
return True
def find_next_prime(start: int, seen: set) -> int:
c = start + 1
while not _is_prime(c) or c in seen:
c += 1
return c
def layout(tc: dict, tp: int) -> dict:
layer_ids = list(tc["engram_layer_ids"])
num_emb = list(tc["engram_num_embeddings"])
max_ngram, n_heads, vocab = tc["engram_max_ngram_size"], tc["engram_n_heads"], tc["engram_vocab_size"]
seen: set = set()
out = {}
for li, layer in enumerate(layer_ids):
sizes = []
for _ in range(max_ngram - 1):
cur = vocab - 1
for _ in range(n_heads):
cur = find_next_prime(cur, seen)
seen.add(cur)
sizes.append(cur)
n_cols = len(sizes)
part = -(-n_cols // tp)
assert sum(sizes) <= num_emb[li], (layer, sum(sizes), num_emb[li])
ranks = []
for r in range(tp):
hs, he = r * part, min(r * part + part, n_cols)
start, end = sum(sizes[:hs]), sum(sizes[:he])
ranks.append({"rank": r, "head_start": hs, "row_start": start, "num_rows": end - start})
out[str(layer)] = {"num_embeddings": num_emb[li], "sum_head_sizes": sum(sizes), "head_sizes": sizes,
"ranks": ranks}
return out
if __name__ == "__main__":
ap = argparse.ArgumentParser()
ap.add_argument("config")
ap.add_argument("--tp", type=int, default=2)
a = ap.parse_args()
cfg = json.load(open(a.config))
print(json.dumps(layout(cfg.get("text_config", cfg), a.tp), indent=1))
|