File size: 1,233 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 | #!/usr/bin/env python3
"""Run the calibration corpus through the live 1M model with the
VLLM_ACT_CAPTURE_DIR hook: per layer, ~24k routed tokens of (x, topk).
Reuses the corpus builder from collect_expert_stats_v2."""
import os
import sys
os.environ["VLLM_ACT_CAPTURE_DIR"] = "/data/glm52-acts"
os.environ.setdefault("VLLM_PP_LAYER_PARTITION", "21,19,19,19")
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from collect_expert_stats_v2 import build_corpus # noqa: E402
def main():
from vllm import LLM, SamplingParams
from vllm.inputs import TokensPrompt
llm = LLM(
model="/data/glm52",
pipeline_parallel_size=4,
gpu_memory_utilization=0.509,
kv_cache_dtype="fp8_ds_mla",
max_model_len=8192,
max_num_seqs=2,
max_num_batched_tokens=2048,
enforce_eager=True,
)
tok = llm.get_tokenizer()
texts = build_corpus(tok)
prompts = [TokensPrompt(prompt_token_ids=tok.encode(t)[:7900])
for t in texts]
print(f"{len(prompts)} calibration prompts")
llm.generate(prompts, SamplingParams(max_tokens=1, temperature=0.0))
print("generation done; check /data/glm52-acts")
if __name__ == "__main__":
main()
|