File size: 4,810 Bytes
fdc6474 | 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 | #!/usr/bin/env python3
"""Collect per-layer expert routing histograms from the hybrid GLM-5.2.
Uses vLLM's enable_return_routed_experts to get [seq_len, n_moe_layers,
topk] expert ids per request, over a mixed local calibration corpus
(code + docs + synthetic instructions). Prefill-only (max_tokens=1).
Output: /data/glm52-expert-stats.npz with counts [n_layers, 256].
"""
import glob
import json
import os
import random
import numpy as np
MODEL = "/data/glm52"
OUT = "/data/glm52-expert-stats.npz"
N_EXPERTS = 256
PROMPT_TOKENS = 6000 # approx, chars/4
N_PROMPTS = 192
def build_corpus() -> list[str]:
random.seed(7)
texts = []
# code: vllm sources
py = sorted(glob.glob("/home/coder/git/glm52/vllm/vllm/**/*.py", recursive=True))
random.shuffle(py)
# prose/tech: docs
md = sorted(glob.glob("/home/coder/git/glm52/vllm/docs/**/*.md", recursive=True))
random.shuffle(md)
def chunks(paths, n):
out, buf = [], ""
for p in paths:
try:
buf += open(p, errors="ignore").read() + "\n\n"
except OSError:
continue
while len(buf) >= PROMPT_TOKENS * 4:
out.append(buf[: PROMPT_TOKENS * 4])
buf = buf[PROMPT_TOKENS * 4:]
if len(out) >= n:
return out
if buf:
out.append(buf)
return out[:n]
texts += chunks(py, N_PROMPTS // 2) # 50% code
texts += chunks(md, N_PROMPTS // 4) # 25% docs
# 25% synthetic instruction/chat/reasoning
topics = ["quantum computing", "the French Revolution", "sourdough bread",
"distributed databases", "protein folding", "jazz harmony",
"supply chains", "volcanoes", "Rust lifetimes", "photosynthesis",
"medieval trade routes", "black holes"]
templates = [
"Explain {t} to a beginner, covering the key concepts step by step, "
"common misconceptions, practical examples, and finally an advanced "
"summary with open research questions. Be very detailed.\n\n",
"Write a detailed technical design document about building a system "
"related to {t}: requirements, architecture, tradeoffs, testing plan, "
"rollout strategy, and failure modes.\n\n",
"You are a helpful assistant. The user asks a long multi-part "
"question about {t}. Answer each part with careful reasoning:\n"
"1) history 2) fundamentals 3) state of the art 4) critiques "
"5) future directions.\n\n",
]
synth = []
while len(synth) < N_PROMPTS // 4:
t = random.choice(topics)
body = random.choice(templates).format(t=t)
synth.append((body * 40)[: PROMPT_TOKENS * 4])
texts += synth
random.shuffle(texts)
return texts
def main():
os.environ.setdefault("VLLM_PP_LAYER_PARTITION", "19,20,21,18")
stats_dir = "/data/glm52-expert-stats"
os.environ["VLLM_HYBRID_EXPERT_STATS"] = stats_dir
from vllm import LLM, SamplingParams
llm = LLM(
model=MODEL,
pipeline_parallel_size=4,
gpu_memory_utilization=0.51,
kv_cache_dtype="fp8_ds_mla",
max_model_len=8192,
max_num_seqs=2,
max_num_batched_tokens=4096,
enforce_eager=True,
)
from vllm.inputs import TokensPrompt
texts = build_corpus()
tok = llm.get_tokenizer()
prompts = [
TokensPrompt(prompt_token_ids=tok.encode(t)[:7900]) for t in texts
]
print(f"{len(prompts)} calibration prompts")
sp = SamplingParams(max_tokens=1, temperature=0.0)
outs = llm.generate(prompts, sp)
total_tokens = sum(len(o.prompt_token_ids) for o in outs)
del llm
# AqlmMoEMethod on each PP worker dumped layer_N.npy files
files = sorted(glob.glob(os.path.join(stats_dir, "layer_*.npy")))
assert files, f"no stats dumped in {stats_dir}"
layer_ids = sorted(int(os.path.basename(f)[6:-4]) for f in files)
counts = np.stack(
[np.load(os.path.join(stats_dir, f"layer_{li}.npy")) for li in layer_ids]
)
np.savez(OUT, counts=counts, layer_ids=np.array(layer_ids),
total_tokens=total_tokens)
# quick skew report
frac = counts / counts.sum(axis=1, keepdims=True).clip(min=1)
sorted_frac = np.sort(frac, axis=1)[:, ::-1]
top32 = sorted_frac[:, :32].sum(axis=1)
top64 = sorted_frac[:, :64].sum(axis=1)
print(f"tokens: {total_tokens}, layers: {counts.shape[0]}")
print(f"routing mass in top-32 experts: mean {top32.mean():.3f} "
f"min {top32.min():.3f} max {top32.max():.3f}")
print(f"routing mass in top-64 experts: mean {top64.mean():.3f} "
f"min {top64.min():.3f} max {top64.max():.3f}")
print("saved", OUT)
if __name__ == "__main__":
main()
|