| |
| """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)) |
|
|
